def submit(self, fun, *args, **kwargs):
"""
Submit a job to the pool.
kwargs may contain __cfut_options which currently should look like:
{
"output_pickle_path": str
}
output_pickle_path defines where the pickled result should be stored.
That file will not be removed after the job has finished.
"""
fut = self.create_enriched_future()
workerid = random_string()
should_keep_output = False
if "__cfut_options" in kwargs:
should_keep_output = True
output_pickle_path = kwargs["__cfut_options"]["output_pickle_path"]
del kwargs["__cfut_options"]
else:
output_pickle_path = self.format_outfile_name(
self.cfut_dir, workerid)
self.ensure_not_shutdown()
# Start the job.
serialized_function_info = pickling.dumps(
(fun, args, kwargs, self.meta_data, output_pickle_path))
with open(self.format_infile_name(self.cfut_dir, workerid), "wb") as f:
f.write(serialized_function_info)
self.store_main_path_to_meta_file(workerid)
preliminary_output_pickle_path = with_preliminary_postfix(
output_pickle_path)
if os.path.exists(preliminary_output_pickle_path):
logging.warning(
f"Deleting stale output file at {preliminary_output_pickle_path}..."
)
os.unlink(preliminary_output_pickle_path)
job_name = get_function_name(fun)
jobids_futures, _ = self._start(workerid, job_name=job_name)
# Only a single job was submitted
jobid = jobids_futures[0].result()
if self.debug:
print("job submitted: %i" % jobid, file=sys.stderr)
# Thread will wait for it to finish.
self.wait_thread.waitFor(preliminary_output_pickle_path, jobid)
with self.jobs_lock:
self.jobs[jobid] = (fut, workerid, output_pickle_path,
should_keep_output)
fut.cluster_jobid = jobid
return fut
def submit_text(cls, job, cfut_dir):
"""Submits a Slurm job represented as a job file string. Returns
the job ID.
"""
filename = cls.get_temp_file_path(
cfut_dir, "_temp_slurm{}.sh".format(random_string()))
with open(filename, "w") as f:
f.write(job)
job_id, stderr = chcall("sbatch --parsable {}".format(filename))
os.unlink(filename)
if len(stderr) > 0:
logging.warning(f"Submitting batch job emitted warnings: {stderr}")
return int(job_id)
def submit_text(self, job):
"""Submits a PBS job represented as a job file string. Returns
the job ID.
"""
filename = self.get_temp_file_path(
self.cfut_dir, "_temp_pbs_{}.sh".format(random_string()))
with open(filename, "w") as f:
f.write(job)
jobid_desc, _ = chcall("qsub -V {}".format(filename))
match = re.search("^[0-9]+", jobid_desc.decode("utf-8"))
assert match is not None
jobid = match.group(0)
print("jobid", jobid)
# os.unlink(filename)
return int(jobid)
def map_to_futures(self, fun, allArgs, output_pickle_path_getter=None):
self.ensure_not_shutdown()
allArgs = list(allArgs)
if len(allArgs) == 0:
return []
should_keep_output = output_pickle_path_getter is not None
futs_with_output_paths = []
workerid = random_string()
pickled_function_path = self.get_function_pickle_path(workerid)
self.files_to_clean_up.append(pickled_function_path)
with open(pickled_function_path, "wb") as file:
pickling.dump(fun, file)
self.store_main_path_to_meta_file(workerid)
for index, arg in enumerate(allArgs):
fut = self.create_enriched_future()
workerid_with_index = self.get_workerid_with_index(workerid, index)
if output_pickle_path_getter is None:
output_pickle_path = self.format_outfile_name(
self.cfut_dir, workerid_with_index)
else:
output_pickle_path = output_pickle_path_getter(arg)
preliminary_output_pickle_path = with_preliminary_postfix(
output_pickle_path)
if os.path.exists(preliminary_output_pickle_path):
logging.warning(
f"Deleting stale output file at {preliminary_output_pickle_path}..."
)
os.unlink(preliminary_output_pickle_path)
serialized_function_info = pickling.dumps(
(pickled_function_path, [arg], {}, self.meta_data,
output_pickle_path))
infile_name = self.format_infile_name(self.cfut_dir,
workerid_with_index)
with open(infile_name, "wb") as f:
f.write(serialized_function_info)
futs_with_output_paths.append((fut, output_pickle_path))
with self.jobs_lock:
# Use a separate loop to avoid having to acquire the jobs_lock many times
# or for the full duration of the above loop
for index in range(len(futs_with_output_paths)):
workerid_with_index = self.get_workerid_with_index(
workerid, index)
# Register the job in the jobs array, although the jobid is not known yet.
# Otherwise it might happen that self.jobs becomes empty, but some of the jobs were
# not even submitted yet.
self.jobs[workerid_with_index] = "pending"
job_count = len(allArgs)
job_name = get_function_name(fun)
jobids_futures, job_index_ranges = self._start(workerid, job_count,
job_name)
number_of_batches = len(jobids_futures)
for batch_index, (jobid_future,
(job_index_start, job_index_end)) in enumerate(
zip(jobids_futures, job_index_ranges)):
jobid_future.add_done_callback(
partial(
self.register_jobs,
futs_with_output_paths[job_index_start:job_index_end],
workerid,
should_keep_output,
job_index_start,
f"{batch_index + 1}/{number_of_batches}",
))
return [fut for (fut, _) in futs_with_output_paths]
