Split host list when multiple nodes are allocated

chart/f7t4jhub/files/jupyterhub-config.py

@@ -140,6 +140,7 @@ async def refresh_user(self, user, handler=None):
 {% if memory %}#SBATCH --mem={{`{{memory}}`}}{% endif %}
 {% if gres %}#SBATCH --gres={{`{{gres}}`}}{% endif %}
 {% if nprocs %}#SBATCH --cpus-per-task={{`{{nprocs}}`}}{% endif %}
+{% if nnodes %}#SBATCH --nodes={{`{{nnodes[0]}}`}}{% endif %}
 {% if reservation%}#SBATCH --reservation={{`{{reservation[0]}}`}}{% endif %}
 {% if constraint %}#SBATCH --constraint={{`{{constraint[0]}}`}}{% endif %}
 {% if options %}#SBATCH {{`{{options}}`}}{% endif %}

dockerfiles/Dockerfile

@@ -1,4 +1,4 @@
-FROM ubuntu:22.04
+FROM ubuntu:24.04
 
 ARG __CONDA_ENV__=py311
 
@@ -20,7 +20,7 @@ RUN . /opt/conda/bin/activate && \
 
 RUN . /opt/conda/bin/activate && \
  conda activate $__CONDA_ENV__ && \
- pip install --no-cache jupyterhub==4.1.6 pyfirecrest==2.1.0 SQLAlchemy==1.4.52 oauthenticator==16.3.1
+ pip install --no-cache jupyterhub==4.1.6 pyfirecrest==2.1.0 SQLAlchemy==1.4.52 oauthenticator==16.3.1 python-hostlist==1.23.0
 
 COPY . firecrestspawner
 RUN . /opt/conda/bin/activate && \
@@ -68,4 +68,4 @@ USER juhu
 
 WORKDIR /home/juhu
 
-CMD . /opt/conda/bin/activate && jupyterhub
+CMD . /opt/conda/bin/activate && conda activate ${__CONDA_ENV__} && jupyterhub

firecrestspawner/spawner.py

@@ -12,6 +12,7 @@
 import re
 import sys
 import jupyterhub
+import hostlist
 import firecrest as f7t
 from enum import Enum
 from jinja2 import Template
@@ -236,6 +237,7 @@ async def submit_batch_script(self):
  self.job = await client.submit(self.host,
  script_str=script,
  env_vars=job_env)
+ self.log.debug(f"[client.submit] {self.job}")
  self.job_id = str(self.job['jobid'])
  self.log.info(f'Job {self.job_id} submitted')
  except Exception as e:
@@ -259,8 +261,11 @@ async def query_job_status(self):
  client = await self.get_firecrest_client()
  self.log.debug('firecREST: Polling job')
  poll_result = await client.poll(self.host, [self.job_id])
+ self.log.debug(f"[client.poll] [query_job_status] {poll_result}")
  state = poll_result[0]['state']
- host = poll_result[0]['nodelist']
+ nodelist = hostlist.expand_hostlist(poll_result[0]['nodelist'])
+ # when PENDING nodelist is []
+ host = nodelist[0] if len(nodelist) > 0 else ""
  # `job_status` must keep the format used in the original
  # batchspawner since it will be later parsed with
  # regular expressions
@@ -284,7 +289,8 @@ async def cancel_batch_job(self):
  self.log.info(f"Cancelling job {self.job_id}")
  client = await self.get_firecrest_client()
  self.log.info('firecREST: Canceling job')
- await client.cancel(self.host, self.job_id)
+ cancel_result = await client.cancel(self.host, self.job_id)
+ self.log.debug(f"[client.cancel] {cancel_result}")
 
  def load_state(self, state):
  """load `job_id` from state"""
@@ -498,9 +504,8 @@ async def state_gethost(self):
 
  client = await self.get_firecrest_client()
  poll_result = await client.poll(self.host, [self.job_id])
- # FIXME: this is expecting `nodelist` to be only a single
- # node. Fix it so it can work with multiple nodes.
- host = poll_result[0]['nodelist']
+ self.log.debug(f"[client.poll] [state_gethost] {poll_result}")
+ host = hostlist.expand_hostlist(poll_result[0]['nodelist'])[0]
  return self.node_name_template.format(host)
 
 

requirements.txt

@@ -2,3 +2,4 @@ jupyterhub==4.1.6
 pyfirecrest==2.1.0
 SQLAlchemy==1.4.52
 oauthenticator==16.3.1
+python-hostlist==1.23.0