matensemble 0.3.7__tar.gz → 0.3.9__tar.gz

{matensemble-0.3.7 → matensemble-0.3.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: matensemble
-Version: 0.3.7
+Version: 0.3.9
 Summary: An adaptive and highly asynchronous ensemble simulation workflow manager MatEnsemble (https://github.com/Q-CAD/MatEnsemble) built jointly on top of the hierarchical graph based scheduler FLUX and concurrent-futures infrastructure of python
 Author: Soumendu Bagchi, Kaleb Duchesneau
 Author-email: Soumendu Bagchi <soumendubagchi@gmail.com>, Kaleb Duchesneau <kalebduchesneau@gmail.com>

{matensemble-0.3.7 → matensemble-0.3.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "matensemble"
-version = "0.3.7"
+version = "0.3.9"
 description = "An adaptive and highly asynchronous ensemble simulation workflow manager MatEnsemble (https://github.com/Q-CAD/MatEnsemble) built jointly on top of the hierarchical graph based scheduler FLUX and concurrent-futures infrastructure of python"
 readme = "README.md"
 license-files = ["LICENSE"]

{matensemble-0.3.7 → matensemble-0.3.9}/src/matensemble/chore.py

@@ -53,6 +53,7 @@ class Chore:
         deps: tuple[str, ...] = (),
         args: tuple = (),
         kwargs: dict | None = None,
+        nnodes: int | None = None,
     ) -> None:
         """
         The constructor for a :obj:`Chore`
@@ -84,6 +85,12 @@ class Chore:
             The arguments to give the function if type is PYTHON
         kwargs : dict
             The key-word arguments to give the function if flaovr is PYTHON
+        nnodes : int, optional
+            When set, this chore will be scheduled via ``per_resource`` and will
+            occupy *nnodes* whole nodes (all cores and all GPUs on each node).
+            The manager uses this to compute the true resource footprint instead
+            of ``num_tasks * cores_per_task`` / ``num_tasks * gpus_per_task``.
+            Leave as ``None`` for normal ``from_command`` chores.
         """
 
         self.id = id
@@ -100,6 +107,7 @@ class Chore:
         self.deps = deps
         self.args = args
         self.kwargs = {} if kwargs is None else kwargs
+        self.nnodes = nnodes
 
     def graph(self) -> nx.DiGraph:
         return nx.DiGraph()
@@ -121,6 +129,7 @@ class Chore:
             "deps": list(self.deps),
             "args": _json_safe(self.args),
             "kwargs": _json_safe(self.kwargs),
+            "nnodes": self.nnodes,
         }
 
     def _write_metadata(self) -> None:

{matensemble-0.3.7 → matensemble-0.3.9}/src/matensemble/dynopro/ensemble.py

@@ -7,6 +7,7 @@ from networkx import write_network_text
 from matensemble.manager import FluxManager
 from matensemble.chore import Chore, ChoreType
 from matensemble.model import Resources
+
 # from matensemble.dynopro.driver import online_dynamics
 
 
@@ -73,6 +74,7 @@ class EnsembleDynamicsRunner:
                     num_tasks=self.tasks_per_job,
                     cores_per_task=self.cores_per_task,
                     gpus_per_task=self.gpus_per_task,
+                    mpi=True,
                 )
                 workdir = outdir / chore_id
 
@@ -85,6 +87,7 @@ class EnsembleDynamicsRunner:
                         chore_type=ChoreType.EXECUTABLE,
                         resources=resources,
                         workdir=workdir,
+                        nnodes=self.nnodes,
                     )
                 )
 

{matensemble-0.3.7 → matensemble-0.3.9}/src/matensemble/fluxlet.py

@@ -26,7 +26,7 @@ class Fluxlet:
         handle: flux.Flux,
     ) -> None:
         self.handle = handle
-        self.gpus_per_node = self.get_gpus_per_node()
+        self.num_nodes, self.gpus_per_node = self.get_gpus_per_node()
 
     def get_gpus_per_node(self) -> tuple[int, int]:
         """
@@ -50,7 +50,6 @@ class Fluxlet:
         chore: Chore,
         set_cpu_affinity: bool | None = None,
         set_gpu_affinity: bool | None = None,
-        nnodes: int | None = None,
         dynopro: bool | None = None,
     ) -> flux.job.FluxExecutorFuture:
         """
@@ -81,11 +80,11 @@ class Fluxlet:
         if dynopro:
             jobspec = flux.job.JobspecV1.per_resource(
                 chore.command,
-                ncores=chore.resources.num_tasks,
-                nnodes=nnodes,
+                nnodes=chore.nnodes,
                 gpus_per_node=self.gpus_per_node,
-                per_resource_type="core",
+                per_resource_type="node",
                 per_resource_count=1,
+                exclusive=True,  # dynopro needs whole nodes
             )
 
             chore.workdir.mkdir(parents=True, exist_ok=True)
@@ -106,8 +105,6 @@ class Fluxlet:
             base_env["SLURM_GPUS_PER_NODE"] = str(self.gpus_per_node)
             jobspec.environment = base_env
 
-            fut = executor.submit(jobspec)
-
             fut = executor.submit(jobspec)
             fut.chore_id = chore.id
             fut.chore_obj = chore

{matensemble-0.3.7 → matensemble-0.3.9}/src/matensemble/manager.py

@@ -272,6 +272,26 @@ class FluxManager:
         gpus_per_node = total_gpus // nnodes
         return nnodes, cores_per_node, gpus_per_node
 
+    def _chore_resource_footprint(self, chore: Chore) -> tuple[int, int]:
+        """
+        Return ``(needed_cores, needed_gpus)`` for a chore.
+
+        For whole-node (dynopro) chores — those with ``chore.nnodes`` set — the
+        footprint is ``nnodes * cores_per_node`` and ``nnodes * gpus_per_node``
+        because ``per_resource`` allocates entire nodes and every core and GPU on
+        them becomes unavailable.  For ordinary chores the footprint is the
+        familiar ``num_tasks * cores_per_task`` / ``num_tasks * gpus_per_task``.
+        """
+
+        if getattr(chore, "nnodes", None) is not None:
+            needed_cores = chore.nnodes * self._cores_per_node
+            needed_gpus = chore.nnodes * self._gpus_per_node
+        else:
+            needed_cores = chore.resources.num_tasks * chore.resources.cores_per_task
+            needed_gpus = chore.resources.num_tasks * chore.resources.gpus_per_task
+
+        return needed_cores, needed_gpus
+
     def _chore_fits_allocation(self, chore: Chore) -> bool:
         """
         Checks whether the given chore is too big to be submitted
@@ -282,8 +302,7 @@ class FluxManager:
             The :obj:`Chore` to check if it will fit in the allocation
         """
 
-        needed_cores = chore.resources.num_tasks * chore.resources.cores_per_task
-        needed_gpus = chore.resources.num_tasks * chore.resources.gpus_per_task
+        needed_cores, needed_gpus = self._chore_resource_footprint(chore)
 
         total_cores = self._nnodes_on_allocation * self._cores_per_node
         total_gpus = self._nnodes_on_allocation * self._gpus_per_node
@@ -334,8 +353,7 @@ class FluxManager:
         Checks to see if there are enough resources to submit the given :obj:`Chore`
         """
 
-        needed_cores = chore.resources.num_tasks * chore.resources.cores_per_task
-        needed_gpus = chore.resources.num_tasks * chore.resources.gpus_per_task
+        needed_cores, needed_gpus = self._chore_resource_footprint(chore)
         return self._free_cores >= needed_cores and self._free_gpus >= needed_gpus
 
     def _has_failed(self, chore_id: str) -> bool:
@@ -399,7 +417,9 @@ class FluxManager:
 
             self._fail_dependents(dep_id)
 
-    def _submit_one(self, chore_id: str, buffer_time: float) -> None:
+    def _submit_one(
+        self, chore_id: str, buffer_time: float, dynopro: bool = False
+    ) -> None:
         """
         Submits a :obj:`Chore` and does book-keeping all the queues and resources
         count
@@ -418,7 +438,7 @@ class FluxManager:
                 chore,
                 set_cpu_affinity=self._set_cpu_affinity,
                 set_gpu_affinity=self._set_gpu_affinity,
-                nnodes=None,
+                dynopro=dynopro,
             )
         except Exception as e:
             self._logger.exception("CHORE SUBMIT FAILED: chore=%s", chore_id)
@@ -441,7 +461,9 @@ class FluxManager:
         self._free_gpus -= chore.resources.num_tasks * chore.resources.gpus_per_task
         time.sleep(buffer_time)
 
-    def _submit_until_ooresources(self, buffer_time: float) -> bool:
+    def _submit_until_ooresources(
+        self, buffer_time: float, dynopro: bool = False
+    ) -> bool:
         """
         Submit as many chores as possible until out-of-resources
 
@@ -459,7 +481,7 @@ class FluxManager:
             chore = self._chores_by_id[chore_id]
 
             if self._can_submit_now(chore):
-                self._submit_one(chore_id, buffer_time)
+                self._submit_one(chore_id, buffer_time, dynopro=dynopro)
                 submitted_any = True
             else:
                 deferred.append(chore_id)
@@ -636,7 +658,7 @@ class FluxManager:
             )
             while not done:
                 self._check_resources()
-                self._submit_until_ooresources(buffer_time=buffer_time)
+                self._submit_until_ooresources(buffer_time=buffer_time, dynopro=dynopro)
                 proc_strat.process_futures(buffer_time=buffer_time)
 
                 done = (

{matensemble-0.3.7 → matensemble-0.3.9}/src/matensemble/strategy.py

@@ -11,6 +11,7 @@ from pathlib import Path
 
 from abc import ABC, abstractmethod
 
+from matensemble import dynopro
 from matensemble.model import OutputReference
 
 
@@ -127,7 +128,10 @@ class AdaptiveStrategy(FutureProcessingStrategy):
                     self.manager._blocked.discard(dep_id)
 
             # adaptively submit another chore
-            self.manager._submit_until_ooresources(buffer_time=buffer_time)
+            self.manager._submit_until_ooresources(
+                buffer_time=buffer_time,
+                dynopro=getattr(self.manager, "_dynopro", False),
+            )
 
             if self.manager._write_restart_freq and (
                 len(self.manager._completed_chores) % self.manager._write_restart_freq
@@ -304,7 +308,7 @@ class UserStrategy(FutureProcessingStrategy):
                     self.manager._blocked.discard(dep_id)
 
             # --- Processing the chore and spawning the new one ---
-            if self.proc_chore == chore_name:
+            if chore_name == self.proc_chore:
                 try:
                     # Trust boundary: result.pickle is written by matensemble.runtime_worker
                     # in this workflow's chore workdir only—do not load pickles from
@@ -339,6 +343,12 @@ class UserStrategy(FutureProcessingStrategy):
                                 f"bolo_match={chore_name} | due the following Exception ->\n{e}"
                             )
 
+            # adaptively submit another chore
+            self.manager._submit_until_ooresources(
+                buffer_time=buffer_time,
+                dynopro=getattr(self.manager, "_dynopro", False),
+            )
+
             if self.manager._write_restart_freq and (
                 len(self.manager._completed_chores) % self.manager._write_restart_freq
                 == 0