def create_queue(): driver = Driver(driver_type=QueueDriverEnum.LOCAL_DRIVER, max_running=5) job_queue = JobQueue(driver) with open(dummy_config["job_script"], "w") as f: f.write("#!/usr/bin/env python\n"\ "with open('STATUS', 'w') as f:"\ " f.write('finished successfully')"\ "\n") os.chmod(dummy_config["job_script"], stat.S_IRWXU | stat.S_IRWXO | stat.S_IRWXG) for i in range(10): os.mkdir(dummy_config["run_path"].format(i)) job = JobQueueNode( job_script=dummy_config["job_script"], job_name=dummy_config["job_name"].format(i), run_path=os.path.realpath(dummy_config["run_path"].format(i)), num_cpu=dummy_config["num_cpu"], status_file=job_queue.status_file, ok_file=job_queue.ok_file, exit_file=job_queue.exit_file, done_callback_function=dummy_config["ok_callback"], exit_callback_function=dummy_config["exit_callback"], callback_arguments=[ { "job_number": i }, os.path.realpath(dummy_config["run_path"].format(i)), ]) job_queue.add_job(job) return job_queue
def create_queue(script, max_submit=2): driver = Driver(driver_type=QueueDriverEnum.LOCAL_DRIVER, max_running=5) job_queue = JobQueue(driver, max_submit=max_submit) with open(dummy_config["job_script"], "w") as f: f.write(script) os.chmod(dummy_config["job_script"], stat.S_IRWXU | stat.S_IRWXO | stat.S_IRWXG) for i in range(10): os.mkdir(dummy_config["run_path"].format(i)) job = JobQueueNode( job_script=dummy_config["job_script"], job_name=dummy_config["job_name"].format(i), run_path=os.path.realpath(dummy_config["run_path"].format(i)), num_cpu=dummy_config["num_cpu"], status_file=job_queue.status_file, ok_file=job_queue.ok_file, exit_file=job_queue.exit_file, done_callback_function=dummy_config["ok_callback"], exit_callback_function=dummy_config["exit_callback"], callback_arguments=[ { "job_number": i }, os.path.realpath(dummy_config["run_path"].format(i)), ]) job_queue.add_job(job) return job_queue
def test_num_cpu_submitted_correctly(self): with TestAreaContext("job_node_test"): os.putenv("PATH", os.getcwd() + ":" + os.getenv("PATH")) driver = Driver(driver_type=QueueDriverEnum.LSF_DRIVER, max_running=1) with open(dummy_config["job_script"], "w") as f: f.write(simple_script) os.chmod(dummy_config["job_script"], stat.S_IRWXU | stat.S_IRWXO | stat.S_IRWXG) with open("bsub", "w") as f: f.write(mock_bsub) os.chmod("bsub", stat.S_IRWXU | stat.S_IRWXO | stat.S_IRWXG) job_id = 0 num_cpus = 4 os.mkdir(dummy_config["run_path"].format(job_id)) job = JobQueueNode( job_script=dummy_config["job_script"], job_name=dummy_config["job_name"].format(job_id), run_path=os.path.realpath( dummy_config["run_path"].format(job_id)), num_cpu=num_cpus, status_file="STATUS", ok_file="OK", exit_file="ERROR", done_callback_function=dummy_config["ok_callback"], exit_callback_function=dummy_config["exit_callback"], callback_arguments=[ { "job_number": job_id }, os.path.realpath(dummy_config["run_path"].format(job_id)), ], ) pool_sema = BoundedSemaphore(value=2) job.run(driver, pool_sema) job.stop() job.wait_for() with open("test.out") as f: bsub_argv = f.read().split() found_cpu_arg = False for arg_i, arg in enumerate(bsub_argv): if arg == "-n": self.assertEqual( bsub_argv[arg_i + 1], str(num_cpus), "num_cpu argument does not match specified number of cpus", ) found_cpu_arg = True self.assertTrue(found_cpu_arg, "num_cpu argument not found")