From 94598859c64c86b80b22f5eda9efe4dd07fb91c9 Mon Sep 17 00:00:00 2001 From: Rafael Sarmiento Date: Fri, 9 Aug 2024 19:26:23 +0200 Subject: [PATCH] loop poll --- firecrestspawner/spawner.py | 26 +++++++++++++++++++++----- 1 file changed, 21 insertions(+), 5 deletions(-) diff --git a/firecrestspawner/spawner.py b/firecrestspawner/spawner.py index e06d238..ff4a1e7 100644 --- a/firecrestspawner/spawner.py +++ b/firecrestspawner/spawner.py @@ -209,6 +209,22 @@ async def get_firecrest_client(self): client.timeout = 30 return client + async def firecrest_poll(self): + """Helper function to poll jobs. + + This is needed in cases that the scheduler is slow updating + its database which could make the result of ``client.poll`` + to be an empty list + """ + client = await self.get_firecrest_client() + poll_result = [] + + while poll_result == []: + poll_result = await client.poll(self.host, [self.job_id]) + print(poll_result) + + return poll_result + async def _get_batch_script(self, **subvars): """Format batch script from vars""" # Could be overridden by subclasses, but mainly useful for testing @@ -258,9 +274,7 @@ async def query_job_status(self): self.log.debug(f'Spawner querying job {self.job_id}') try: - client = await self.get_firecrest_client() - self.log.debug('firecREST: Polling job') - poll_result = await client.poll(self.host, [self.job_id]) + poll_result = await self.firecrest_poll() self.log.debug(f"[client.poll] [query_job_status] {poll_result}") state = poll_result[0]['state'] nodelist = hostlist.expand_hostlist(poll_result[0]['nodelist']) @@ -502,9 +516,11 @@ async def state_gethost(self): if self.custom_state_gethost: return await self.custom_state_gethost(self) - client = await self.get_firecrest_client() - poll_result = await client.poll(self.host, [self.job_id]) + poll_result = await self.firecrest_poll() self.log.debug(f"[client.poll] [state_gethost] {poll_result}") + + # this function is called only when the job has been allocated, + # then ``nodelist`` won't be ``[]`` host = hostlist.expand_hostlist(poll_result[0]['nodelist'])[0] return self.node_name_template.format(host)