| 52 | |
| 53 | |
| 54 | class Client: |
| 55 | def __init__( |
| 56 | self, |
| 57 | server_dir: Optional[GenericPath] = None, |
| 58 | python_env: Optional[PythonEnv] = None, |
| 59 | ): |
| 60 | """ |
| 61 | A client serves as a gateway for submitting jobs and querying information about a running |
| 62 | HyperQueue server. |
| 63 | |
| 64 | :param server_dir: Path to a server directory of a running HyperQueue server. |
| 65 | :param python_env: Python environment which configures Python tasks created by |
| 66 | [`function`](`hyperqueue.job.Job.function`). |
| 67 | """ |
| 68 | server_dir = str(server_dir) if server_dir else None |
| 69 | self.connection = ClientConnection(server_dir) |
| 70 | if python_env is None: |
| 71 | python_env = PythonEnv() |
| 72 | self.python_env = python_env |
| 73 | |
| 74 | def submit(self, job: Job) -> SubmittedJob: |
| 75 | """ |
| 76 | Submit a job into HyperQueue. |
| 77 | |
| 78 | :param job: Job that will be submitted. |
| 79 | """ |
| 80 | job_desc = job._build(self) |
| 81 | task_count = len(job_desc.tasks) |
| 82 | if task_count < 1: |
| 83 | raise Exception("Submitted job must have at least a single task") |
| 84 | |
| 85 | job_id = self.connection.submit_job(job_desc) |
| 86 | logging.info(f"Submitted job {job_id} with {task_count} {pluralize('task', task_count)}") |
| 87 | return SubmittedJob(job=job, id=job_id) |
| 88 | |
| 89 | def wait_for_jobs(self, jobs: Sequence[SubmittedJob], raise_on_error=True) -> bool: |
| 90 | """Returns True if all tasks were successfully finished""" |
| 91 | |
| 92 | job_ids = tuple(job.id for job in jobs) |
| 93 | job_ids_str = ",".join(str(id) for id in job_ids) |
| 94 | if len(jobs) > 1: |
| 95 | job_ids_str = "{" + job_ids_str + "}" |
| 96 | logging.info(f"Waiting for {pluralize('job', len(jobs))} {job_ids_str} to finish") |
| 97 | |
| 98 | callback = create_progress_callback() |
| 99 | |
| 100 | failed_jobs = self.connection.wait_for_jobs(job_ids, callback) |
| 101 | if failed_jobs and raise_on_error: |
| 102 | failed_tasks = self.connection.get_failed_tasks(failed_jobs) |
| 103 | job_map = {job.id: job.job for job in jobs} |
| 104 | raise FailedJobsException(failed_tasks, job_map) |
| 105 | return len(failed_jobs) == 0 |
| 106 | |
| 107 | def get_failed_tasks(self, job: SubmittedJob) -> Dict[TaskId, FailedTaskContext]: |
| 108 | result = self.connection.get_failed_tasks([job.id]) |
| 109 | return result[job.id] |
| 110 | |
| 111 | def forget(self, job: HasJobId): |
no outgoing calls