Merge pull request #29 from eth-cscs/chart-update-

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