libpipe 0.9__tar.gz → 0.10__tar.gz

{libpipe-0.9 → libpipe-0.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: libpipe
-Version: 0.9
+Version: 0.10
 Summary: Pipeline orchestration and node management utilities
 License-Expression: GPL-3.0-or-later
 License-File: LICENSE

{libpipe-0.9 → libpipe-0.10}/libpipe/worker.py

@@ -51,6 +51,10 @@ ALLOC_MAX_REQUESTS = 3
 ALLOC_WORKERS_PER_NODE = 32
 
 
+def now_str():
+    return datetime.datetime.now().isoformat(sep=' ', timespec='milliseconds')
+
+
 def mem_mb(text):
     '''Memory such as 40GB, 40G or 512MB in MB (GB is 1024 MB, as in Slurm).'''
     m = re.fullmatch(r'\s*(\d+(?:\.\d+)?)\s*([KMGT]?)(?:i?B)?\s*', str(text), re.IGNORECASE)
@@ -259,7 +263,13 @@ class Task:
             line = style(line, fg='red')
         print(f'[{self.name}:{host}] {line}')
 
-    async def log_stream(self, stream, host, err=False):
+    def log_event(self, text):
+        """Write a timestamped libpipe event to the log file only."""
+        if self.fd is not None:
+            self.fd.write(f'[libpipe] {now_str()} {text}\n')
+
+    async def log_stream(self, stream, host, err=False, stamp_prefix=None):
+        """Log the lines of stream; those starting with stamp_prefix get the current time."""
         while True:
             try:
                 line = await stream.readline()
@@ -268,6 +278,8 @@ class Task:
                 continue
             if not line:
                 break
+            if stamp_prefix and line.startswith(stamp_prefix):
+                line = f'{now_str()} '.encode() + line
             self.log(line, host, err=err)
 
     def close_log(self):
@@ -618,15 +630,14 @@ class SlurmClient:
             if self.alloc_id:
                 args.append(f'--jobid={self.alloc_id}')
             if self.exact:
-                # --quiet: no "step creation disabled, retrying" while a step waits for its node
-                args += ['--exact', '--quiet']
+                args.append('--exact')
             if task.cpus:
                 args += [f'--cpus-per-task={task.cpus}', f'--mem={task.mem_mb}M']
                 threads = ' '.join(f'{v}={task.cpus}' for v in THREAD_VARS)
                 command = f'export {threads}; {command}'
             job_id, options = '${SLURM_JOB_ID}.${SLURM_STEP_ID}', []
         args += self._time_option(options) + options
-        marker = f'echo "{SLURM_MARKER}{job_id} on $(hostname -s)"'
+        marker = f'echo "{SLURM_MARKER}{job_id} on $(hostname -s) at $(date \'+%F %T.%3N\')"'
         return f'{shlex.join(args)} bash -c {shlex.quote(f"{marker}; {command}")}'
 
     async def _run(self, *args):
@@ -693,26 +704,42 @@ class SlurmClient:
             text = line.decode(errors='replace')
             if task.slurm_job_id is None and text.startswith(SLURM_MARKER):
                 task.slurm_job_id = text[len(SLURM_MARKER):].split()[0]
+                task.start_time = time.time()
             task.log(line, self.host)
 
     async def execute(self, task):
         task.slurm_job_id = None
         if self.alloc_nodes:
             await self.ensure_allocation()
-        cpus = await self._acquire_cpus(task.cpus or 1) if self.steps else 0
+        cpus = 0
+        if self.steps and self.cpus_total is not None:
+            task.log_event(f'waiting for {min(task.cpus or 1, self.cpus_total)} CPUs '
+                           f'({self.cpus_free} of {self.cpus_total} free)')
+            cpus = await self._acquire_cpus(task.cpus or 1)
         try:
             process = await asyncio.create_subprocess_shell(
                 self.srun_command(task), stdout=asyncio.subprocess.PIPE,
                 stderr=asyncio.subprocess.PIPE, preexec_fn=os.setpgrp, env=self.env)
+            task.log_event(f'srun started{self._cpus_state()}')
             task.set_process(Process(self, process))
+            # srun waits for the job or the step to start: running from the marker line
+            task.start_time = None
 
             await asyncio.gather(
                 self._log_stdout(task, process.stdout),
-                task.log_stream(process.stderr, self.host, err=True),
+                task.log_stream(process.stderr, self.host, err=True, stamp_prefix=b'srun: '),
             )
             task.returncode = await process.wait()
+            task.log_event(f'srun exited with code {task.returncode}')
         finally:
             await self._release_cpus(cpus)
+            if cpus:
+                task.log_event(f'{cpus} CPUs released{self._cpus_state()}')
+
+    def _cpus_state(self):
+        if self.cpus_total is None:
+            return ''
+        return f' ({self.cpus_free} of {self.cpus_total} CPUs free)'
 
     async def terminate(self, task, sig=signal.SIGTERM):
         job_id = getattr(task, 'slurm_job_id', None)
@@ -1096,6 +1123,7 @@ class WorkerPool:
                 tasks = self.get_all_tasks()
                 n_done = len(tasks[TaskStatus.SUCCESS])
                 n_todo = len(tasks[TaskStatus.WAITING]) + len(tasks[TaskStatus.RUNNING])
+                self._start_pbar_timer()
                 self.pbar.update(n_done)
 
                 if self.max_time is not None:
@@ -1104,6 +1132,15 @@ class WorkerPool:
         except asyncio.CancelledError:
             pass
 
+    def _start_pbar_timer(self):
+        """Time the pool from the first task start, without the wait for the Slurm jobs."""
+        if self.pbar_timed:
+            return
+        starts = [t.start_time for t in self.tasks if t.start_time is not None]
+        if starts:
+            self.pbar.start_time = datetime.datetime.fromtimestamp(min(starts))
+            self.pbar_timed = True
+
     async def _enforce_max_time(self):
         now = time.time()
         for worker in self.workers:
@@ -1150,6 +1187,7 @@ class WorkerPool:
         self.tasks.sort(key=lambda task: not task.run_on_host)
 
         self.pbar.start(max_value=len(self.tasks))
+        self.pbar_timed = False
 
         # An allocation no larger than the number of tasks
         for client in {w.client for w in self.workers}:

{libpipe-0.9 → libpipe-0.10}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "libpipe"
-version = "0.9"
+version = "0.10"
 description = "Pipeline orchestration and node management utilities"
 readme = "README.md"
 license = "GPL-3.0-or-later"