Implements manager logic on the level of CLI, reading command-line arguments and returning an exit status and standard output.
| 125 | |
| 126 | |
| 127 | class CommandHandler: |
| 128 | """ |
| 129 | Implements manager logic on the level of CLI, reading command-line arguments |
| 130 | and returning an exit status and standard output. |
| 131 | """ |
| 132 | |
| 133 | def __init__(self, adapter: ManagerAdapter): |
| 134 | self.adapter = adapter |
| 135 | |
| 136 | self.job_counter = 0 |
| 137 | # None = job is purposely missing |
| 138 | self.jobs: Dict[str, Optional[JobData]] = {} |
| 139 | self.deleted_jobs = set() |
| 140 | |
| 141 | async def handle_command(self, input: CommandInput) -> CommandOutput: |
| 142 | cmd_type = self.adapter.parse_command_type(input) |
| 143 | if cmd_type == "submit": |
| 144 | return await self.handle_cli_submit(input) |
| 145 | elif cmd_type == "status": |
| 146 | return await self.handle_cli_status(input) |
| 147 | elif cmd_type == "delete": |
| 148 | return await self.handle_cli_delete(input) |
| 149 | raise Exception(f"Input command {input} not handled") |
| 150 | |
| 151 | async def handle_cli_submit(self, input: CommandInput) -> CommandOutput: |
| 152 | job_id = await self.handle_submit() |
| 153 | return self.adapter.format_submit_output(job_id) |
| 154 | |
| 155 | async def handle_submit(self) -> JobId: |
| 156 | # By default, create a new job |
| 157 | job_id = default_job_id(self.job_counter) |
| 158 | self.job_counter += 1 |
| 159 | |
| 160 | # The state of this job could already have been set before manually |
| 161 | if job_id not in self.jobs: |
| 162 | self.jobs[job_id] = JobData.queued() |
| 163 | return job_id |
| 164 | |
| 165 | async def handle_cli_status(self, input: CommandInput) -> CommandOutput: |
| 166 | job_ids = self.adapter.parse_status_job_ids(input) |
| 167 | job_data = await self.handle_status(job_ids) |
| 168 | return self.adapter.format_status_output(job_data) |
| 169 | |
| 170 | async def handle_status(self, job_ids: List[JobId]) -> Dict[JobId, Optional[JobData]]: |
| 171 | return {job_id: self.jobs.get(job_id) for job_id in job_ids} |
| 172 | |
| 173 | async def handle_cli_delete(self, input: CommandInput) -> CommandOutput: |
| 174 | assert len(input.arguments) == 1 |
| 175 | await self.handle_delete(input.arguments[0]) |
| 176 | return response() |
| 177 | |
| 178 | async def handle_delete(self, job_id: JobId): |
| 179 | assert job_id in self.jobs |
| 180 | self.deleted_jobs.add(job_id) |
| 181 | |
| 182 | def add_worker(self, hq_env: HqEnv, allocation_id: str) -> Popen: |
| 183 | self.set_job_data(allocation_id, JobData.running()) |
| 184 | return self.adapter.start_worker(hq_env, allocation_id) |
no outgoing calls