matensemble 0.2.1__tar.gz

matensemble-0.2.1/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2025, Soumendu Bagchi
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

matensemble-0.2.1/PKG-INFO

@@ -0,0 +1,124 @@
+Metadata-Version: 2.4
+Name: matensemble
+Version: 0.2.1
+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>
+License-File: LICENSE
+Requires-Dist: ase>=3.22.0
+Requires-Dist: cloudpickle>=3.1.2
+Requires-Dist: matplotlib>=3.4.0
+Requires-Dist: networkx>=3.6.1
+Requires-Dist: numpy>=1.21.0
+Requires-Dist: ovito>=3.7.0
+Requires-Dist: pandas>=1.3.0
+Requires-Dist: pymatgen>=2022.0.0
+Requires-Dist: scikit-learn>=1.0.0
+Requires-Dist: scipy>=1.7.0
+Requires-Dist: seaborn>=0.11.0
+Requires-Dist: flux-python==0.66.0 ; extra == 'flux'
+Requires-Python: >=3.12
+Provides-Extra: flux
+Description-Content-Type: text/markdown
+
+[![PyPI version](https://badge.fury.io/py/matensemble.svg)](https://pypi.org/project/matensemble/)
+[![Documentation](https://readthedocs.org/projects/matensemble/badge/?version=latest)](https://matensemble.readthedocs.io/en/latest/)
+[![Python](https://img.shields.io/badge/python-3.12%2B-blue)](https://www.python.org/downloads/)
+[![License](https://img.shields.io/badge/License-BSD%203--Clause-blue.svg)](https://opensource.org/licenses/BSD-3-Clause)
+
+<p align="center">
+  <img src="images/Logo-Matensemble.png" alt="MatEnsemble" width="720" />
+</p>
+
+# MatEnsemble
+
+MatEnsemble is a Python library for **high-throughput workflows** on HPC systems. You define a directed acyclic graph (DAG) of tasks—**Python callables** or **executable commands**—and MatEnsemble submits work through **[Flux](https://flux-framework.readthedocs.io/)**, tracks completions, **adapts** scheduling to free CPUs and GPUs, and writes structured logs and per-chore output directories.
+
+An optional in-tree **dynopro** stack supports streaming dynamics and on-the-fly analysis for advanced materials simulation workflows.
+
+## Features
+
+- **DAG-based workflows** with dependencies via deferred return values (`OutputReference`)
+- **Adaptive scheduling** that back-fills the allocation as tasks finish (with a non-adaptive mode when you need it)
+- **Two chore types**: Python chores (remotely unpickled and executed by `matensemble.runtime_worker`) and argv-style **executable** chores
+- **Resource requests**: tasks, cores per task, GPUs per task, optional MPI (`pmi2`) via Flux
+- **Observability**: `status.json`, `matensemble_workflow.log`, per-chore `stdout` / `stderr`, pickle and JSON result artifacts; optional **web dashboard** (FastAPI on port 8000)
+
+<p align="center">
+  <img src="images/Cap_1_adaptive_task_management.png" alt="Adaptive task management" width="620" />
+</p>
+
+<p align="center">
+  <img src="images/Cap_2_dynopro.png" alt="On-the-fly dynamics and analysis" width="620" />
+</p>
+
+## Documentation
+
+Documentation (overview, architecture, tutorials, API reference):
+
+**[matensemble.readthedocs.io](https://matensemble.readthedocs.io/en/latest/)**
+
+## Installation
+
+### Containers (recommended on many clusters)
+
+OCI images are published to GitHub Container Registry, for example:
+
+`ghcr.io/freddude2004/matensemble:baseline-vX.Y.Z`
+
+See the [container packages](https://github.com/FredDude2004/MatEnsemble/pkgs/container/matensemble) and the [Quick start](https://matensemble.readthedocs.io/en/latest/quickstart.html) in the docs for Apptainer/Singularity and site-specific notes.
+
+### Development install
+
+From a clone of this repository:
+
+```bash
+uv sync
+uv sync --group dev   # optional: docs and pytest tooling
+uv run pytest
+```
+
+Or with pip:
+
+```bash
+pip install -e ".[flux]"
+```
+
+Site-specific Conda-style environment files live under `scripts/` (for example `scripts/baseline/environment.yaml`, `scripts/frontier/`, `scripts/perlmuter/`). Align Python with **3.12+** and Flux with your center’s modules.
+
+## Quick example
+
+```python
+from matensemble.pipeline import Pipeline
+
+pipe = Pipeline()
+pipe.exec(command=["/bin/echo", "hello from MatEnsemble"])
+pipe.submit()
+```
+
+For Python chores, dependency graphs, and the required split between an importable **chore module** and a **runner script**, see the [Tutorials](https://matensemble.readthedocs.io/en/latest/tutorials.html).
+
+## Examples in the repository
+
+Illustrative workflows live under [`example_workflows/`](https://github.com/FredDude2004/MatEnsemble/tree/main/example_workflows).
+
+## Requirements and runtime
+
+- A **Flux allocation** (or equivalent) on the machine where you call `Pipeline.submit()`
+- For MPI Python or executable chores: a coherent MPI/Flux setup (e.g. PMI2) as expected by your site
+- Optional: SSH port forwarding if you enable the dashboard on a compute node (see the architecture guide in the docs)
+
+## Related links
+
+- [Flux documentation](https://flux-framework.readthedocs.io/)
+- [Flux Python guide](https://flux-framework.readthedocs.io/projects/flux-core/en/latest/guide/start.html)
+- [Slurm documentation](https://slurm.schedmd.com/documentation.html) (common front-end to batch allocations)
+- [LAMMPS manual](https://docs.lammps.org/Manual.html) (often used alongside ensemble MD workflows)
+
+## Authors
+
+Soumendu Bagchi, Kaleb Duchesneau (see `pyproject.toml` for contact details).
+
+## License
+
+BSD 3-Clause. See [`LICENSE`](LICENSE).

matensemble-0.2.1/README.md

@@ -0,0 +1,101 @@
+[![PyPI version](https://badge.fury.io/py/matensemble.svg)](https://pypi.org/project/matensemble/)
+[![Documentation](https://readthedocs.org/projects/matensemble/badge/?version=latest)](https://matensemble.readthedocs.io/en/latest/)
+[![Python](https://img.shields.io/badge/python-3.12%2B-blue)](https://www.python.org/downloads/)
+[![License](https://img.shields.io/badge/License-BSD%203--Clause-blue.svg)](https://opensource.org/licenses/BSD-3-Clause)
+
+<p align="center">
+  <img src="images/Logo-Matensemble.png" alt="MatEnsemble" width="720" />
+</p>
+
+# MatEnsemble
+
+MatEnsemble is a Python library for **high-throughput workflows** on HPC systems. You define a directed acyclic graph (DAG) of tasks—**Python callables** or **executable commands**—and MatEnsemble submits work through **[Flux](https://flux-framework.readthedocs.io/)**, tracks completions, **adapts** scheduling to free CPUs and GPUs, and writes structured logs and per-chore output directories.
+
+An optional in-tree **dynopro** stack supports streaming dynamics and on-the-fly analysis for advanced materials simulation workflows.
+
+## Features
+
+- **DAG-based workflows** with dependencies via deferred return values (`OutputReference`)
+- **Adaptive scheduling** that back-fills the allocation as tasks finish (with a non-adaptive mode when you need it)
+- **Two chore types**: Python chores (remotely unpickled and executed by `matensemble.runtime_worker`) and argv-style **executable** chores
+- **Resource requests**: tasks, cores per task, GPUs per task, optional MPI (`pmi2`) via Flux
+- **Observability**: `status.json`, `matensemble_workflow.log`, per-chore `stdout` / `stderr`, pickle and JSON result artifacts; optional **web dashboard** (FastAPI on port 8000)
+
+<p align="center">
+  <img src="images/Cap_1_adaptive_task_management.png" alt="Adaptive task management" width="620" />
+</p>
+
+<p align="center">
+  <img src="images/Cap_2_dynopro.png" alt="On-the-fly dynamics and analysis" width="620" />
+</p>
+
+## Documentation
+
+Documentation (overview, architecture, tutorials, API reference):
+
+**[matensemble.readthedocs.io](https://matensemble.readthedocs.io/en/latest/)**
+
+## Installation
+
+### Containers (recommended on many clusters)
+
+OCI images are published to GitHub Container Registry, for example:
+
+`ghcr.io/freddude2004/matensemble:baseline-vX.Y.Z`
+
+See the [container packages](https://github.com/FredDude2004/MatEnsemble/pkgs/container/matensemble) and the [Quick start](https://matensemble.readthedocs.io/en/latest/quickstart.html) in the docs for Apptainer/Singularity and site-specific notes.
+
+### Development install
+
+From a clone of this repository:
+
+```bash
+uv sync
+uv sync --group dev   # optional: docs and pytest tooling
+uv run pytest
+```
+
+Or with pip:
+
+```bash
+pip install -e ".[flux]"
+```
+
+Site-specific Conda-style environment files live under `scripts/` (for example `scripts/baseline/environment.yaml`, `scripts/frontier/`, `scripts/perlmuter/`). Align Python with **3.12+** and Flux with your center’s modules.
+
+## Quick example
+
+```python
+from matensemble.pipeline import Pipeline
+
+pipe = Pipeline()
+pipe.exec(command=["/bin/echo", "hello from MatEnsemble"])
+pipe.submit()
+```
+
+For Python chores, dependency graphs, and the required split between an importable **chore module** and a **runner script**, see the [Tutorials](https://matensemble.readthedocs.io/en/latest/tutorials.html).
+
+## Examples in the repository
+
+Illustrative workflows live under [`example_workflows/`](https://github.com/FredDude2004/MatEnsemble/tree/main/example_workflows).
+
+## Requirements and runtime
+
+- A **Flux allocation** (or equivalent) on the machine where you call `Pipeline.submit()`
+- For MPI Python or executable chores: a coherent MPI/Flux setup (e.g. PMI2) as expected by your site
+- Optional: SSH port forwarding if you enable the dashboard on a compute node (see the architecture guide in the docs)
+
+## Related links
+
+- [Flux documentation](https://flux-framework.readthedocs.io/)
+- [Flux Python guide](https://flux-framework.readthedocs.io/projects/flux-core/en/latest/guide/start.html)
+- [Slurm documentation](https://slurm.schedmd.com/documentation.html) (common front-end to batch allocations)
+- [LAMMPS manual](https://docs.lammps.org/Manual.html) (often used alongside ensemble MD workflows)
+
+## Authors
+
+Soumendu Bagchi, Kaleb Duchesneau (see `pyproject.toml` for contact details).
+
+## License
+
+BSD 3-Clause. See [`LICENSE`](LICENSE).

matensemble-0.2.1/pyproject.toml

@@ -0,0 +1,48 @@
+[project]
+name = "matensemble"
+version = "0.2.1"
+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"]
+authors = [
+    { name = "Soumendu Bagchi", email = "soumendubagchi@gmail.com" },
+    { name = "Kaleb Duchesneau", email = "kalebduchesneau@gmail.com" },
+]
+requires-python = ">=3.12"
+dependencies = [
+    "ase>=3.22.0",
+    "cloudpickle>=3.1.2",
+    "matplotlib>=3.4.0",
+    "networkx>=3.6.1",
+    "numpy>=1.21.0",
+    "ovito>=3.7.0",
+    "pandas>=1.3.0",
+    "pymatgen>=2022.0.0",
+    "scikit-learn>=1.0.0",
+    "scipy>=1.7.0",
+    "seaborn>=0.11.0",
+]
+
+[project.optional-dependencies]
+flux = [
+    "flux-python==0.66.0",
+]
+
+[build-system]
+requires = ["uv_build>=0.10.0,<0.11.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "myst-parser>=5.0.0",
+    "pytest>=9.0.2",
+    "sphinx>=9.1.0",
+    "sphinx-autodoc-typehints>=3.6.2",
+    "sphinx-rtd-theme>=3.1.0",
+]
+
+[tool.uv.workspace]
+members = [
+    "src/mcp",
+    "src/mcp_matensemble",
+]

matensemble-0.2.1/src/matensemble/.python-version

@@ -0,0 +1 @@
+3.12

matensemble-0.2.1/src/matensemble/TODO.md

@@ -0,0 +1,208 @@
+# === TODO ===
+
+## --- Refactor MatEnsemble ---
+- [x] Create strategy base class
+- [x] Implement the stategies
+- [x] Create strategy base class for processing futures
+- [x] Implement strategies for processing futures
+- [x] Refactor matflux.py and matfluxGen.py to be more modular and in one manager.py
+- [x] Test matflux/matfluxGen refactor make sure it works before doing anything else
+- [x] Fix problems and test again
+
+### --- Problems ---
+- [x] Use *ONE* executor in the manager super loop instead of spawning new ones each time
+- [x] Make sure future objects have proper fields appended at creation (task_ or task + job_spec)
+- [x] Move writing of restart files into the FutureProcessingStrategy implementations
+- [x] Make sure you remove the finished future rather than popleft in FutureProcessingStrategy implementations
+
+## --- NOTE: Refactored code runs way slower ---
+- [x] Fix problems causing slowdown and test again
+
+### --- More Problems ---
+- [x] Make tests consistent so that we have an apples to apples comparison
+- [x] Remove extra logging and RPC calls to limit traffic
+- [x] Update resources calls to update in place in submit_until_ooresources()
+
+- [x] Test matensemble again until it is working as before
+
+  **--- Got it working as before ---**
+- [x] Update logging to be more industry standard
+- [x] Refactor Fluxlet to remove global side effects
+- [x] Add type annotations back to strategies
+- [x] Document all of the code vigorously
+    - [x] Document manager.py
+    - [x] Document fluxlet.py
+    - [x] Document strategies/*
+- [x] Remove all TODOs and HACKs
+
+- [x] Update the documentation and make sure it has all of the strategies
+- [x] Make a script to build the documentation
+- [x] Remove all artifacts from the repository
+
+## --- Add Testing ---
+- [x] Make sure the simple hello world tests work
+- [x] Figure out what is going on with the GPU tasks
+- [x] Make some tests that have failures to make sure the failed tasks get logged appropriately
+- [x] NOTE: Come back later -- Unit tests | Integration Tests --
+
+## --- Find Solution For Distribution ---
+- [x] Turn Matensemble into a uv project
+- [x] Build the initial Apptainer container
+    - [x] Create the matensemble.def file
+    ### --- Def File Spec ---
+    - The file should be off of a frontier base image use rocky linux version
+    - Install build dependecies of flux-core
+    - Build flux-core from source
+    - Install build dependencies for flux-sched
+    - Build flux-sched from source
+    - Export all variables
+    - Install matensemble
+- [x] Test the apptainer container
+
+## --- Build Base Images ---
+- [x] Build Base Image for Baseline
+- [x] Build Base Image for Frontier
+- [x] Build Base Image for Perlmutter
+- [x] Push images to GitHub Container Registry
+- [x] Build MatEnsemble Images with each base image
+- [x] Test Images on each respective system
+- [x] NOTE: Come back later -- Test Perlmutter Image
+- [x] Create Perlmutter image from new base image [Neil's Containerfiles](https://github.com/namehta4/Containerfiles/blob/main/Base/GPU/Dockerfile)
+
+## --- Setup GitHub Actions ---
+- [x] Setup Matrix build action to build MatEnsemble images for baseline, frontier, and perlmutter and push them to ghcr
+- [x] Setup action to build with uv and publish with uv
+- [x] Setup action to build docs and publish them
+
+## --- Test CI/CD ---
+- [x] Make small change to MatEnsemble and Docs
+- [x] Push to main see if dev builds succeed
+- [x] Run release.sh script to see if releases happen properly
+
+## --- Updated/Better UX ---
+- [x] Refactor to be built around Task/Job Objects
+- [x] Allow users to decorate python functions to create TaskSpec's
+    - [x] Allow functions to depend on other functions
+    - [x] Topologically sort all of the Jobs based on dependencies
+- [x] Write worker runtime that flux can target and call user defined functions
+- [x] Write the Job objects specification to a file in their direcotry
+
+## --- Status Dashboard ---
+- [x] Update the Pipeline.run() method to have a dashboard flag
+- [x] Add logic to launch the dashboard when the user runs the workflow
+
+## --- Science Example ---
+- [x] Test the science example that Soumendu provided
+- [x] Update version of LAMMPS
+- [x] Fix Bug with jobspec.env -> jobspec.environment
+- [x] Test it again
+
+## --- Polish Everything ---
+- [x] Update all the documentation
+- [x] Update the example workflows
+- [x] Provide tutorials for how to run the example workflows
+- [x] Change name of 'Job' to 'Chore'
+- [x] Change name of 'Pipeline' to something else
+- [x] Make ChoreType.PYTHON have the ability to be defined in the runner script
+
+## --- Fix Containers ---
+- [x] Install latest version of lammps in frontier images
+- [x] Test MPI problem with MatEnsemble in Frontier Images
+- [x] Test Science Example with latest MPICH install
+- [x] Ping neil to ask about MPICH in image
+
+## --- Finish presentation ---
+- [x] Conda environment might make this very simple
+- [x] Create a jupyter notebook and screen record it
+- [x] Place that at the end of the presentation
+
+## MatEnsemble Fixes
+- [x] Add the ability to print the results of the OutputReference objects
+- [x] Give the user the ability to define workflows in a single file
+- [x] Make the log updates threaded
+- [x] Implement the restart files
+- [x] Bring back the terminal view the log command
+- [x] Dynopro fix thingy (read flux docs JobspecV1.from_command vs. JobspecV1.per_resource)
+- [x] Test all the fixes with a simple dynopro example and print an
+      OutputReference and watch the logs and make a single file workflow and
+
+## --- Test Perlmutter Container ---
+- [x] Give them a test with the current command that you have been running
+- [x] If that doesn't work break it down into smaller pieces
+  ## --- Smaller pieces ---
+  - [x] Maybe start with an nvidia image rather than Neil's image
+  - [ ] Make sure that flux works in the container
+  - [ ] Create a container that just has flux and have some different tests for that
+  - [ ] Create a container that has just MPI and test that make sure it works
+  - [ ] Combine flux and MPI and see if that works
+  - [ ] Create a container that has lammps and make sure that that is working
+  - [ ] Combine all the pieces
+
+## --- Test Frontier Apptainer container ---
+- [x] Need lots more help here
+- [x] Make test that is very small first
+- [x] Test new container against the running examples
+
+## --- Clean some things up ---
+- [x] Make ticket for neil
+- [x] Update frontier Dockerfile
+- [x] Update README.md
+- [x] Create multi architecture builds for baseline image
+- [x] Run release script
+
+## --- Create new strategy to enable autonomous workflows ---
+- [x] Update the OutputReference objects to have the ability to get the results
+- [x] Create a method in the Pipeline to be able to get the results of all chores
+- [x] Make a strategy that can take in chore and does processing which spawns a new chore
+- [x] Figure out how to spawn a new chore
+- [x] Change to only use cloudpickle and only pickle the actual function once
+- [x] Change the chore objects to not store the function
+- [x] Change the chore objects to reference a function in the registry
+- [x] Change runtime worker to load function from registry and call it with args/kwargs
+- [x] Add a set of OutputReference objects in the pipeline
+- [x] Connect the added chores back to the pipeline somehow
+- [x] Add a method in pipeline where you can get the results of all of your chores
+- [x] Ping Neil about MPICH
+- [x] Convert Scaffold into PowerPoint Presentation
+- [x] Email coordinatior about length of presentation and audience
+
+## === AFTER EVERYTHING ABOVE IS DONE AND STABLE ===
+
+## --- Model Context Protocol ---
+- [ ] MCP implementation
+- [ ] Map out the Tool and Resources
+    ### Implement the Resources
+    - [ ] Resource to Fetch ALL Docs
+    - [ ] Resource to Fetch Relavant Source Code
+    - [ ] Resource to Fetch Examples General or system dependent
+    ### Implement the Tools
+    - [ ] Tool to create a directory for the workflow
+    - [ ] Tool to write a file in that directory
+    - [ ] Tool to delete a file in that directory
+    - [ ] Tool to create a workflow
+    - [ ] Tool to verify a workflow
+    - [ ] Tool to create a batch script
+    - [ ] Tool to setup container env
+    - [ ] Tool to submit a batch script
+    ### Implement the Prompts
+    - [ ] ???
+- [ ] Test the server locally
+- [ ] Test the server on an HPC cluster
+- [ ] Create documentation for setting it up
+
+## --- Create first draft for JOSS ---
+- [ ] Read some example papers
+- [ ] Create draft and show Dr. Bagchi
+- [ ] Polish the repository to be ready for review
+- [ ] Make sure that the tests work
+- [ ] Make sure that the example workflows work correctly
+- [ ] Make sure that they can easily test the code and
+- [ ] Create a conda package that they can easily test the code without having
+      to compile flux and flux-sched themselves
+
+## --- Reading List ---
+- [ ] [Agentic Orchestration of HPC Applications](https://vsoch.github.io/assets/posts/agentic-orchestration-hpc-workloads-cloud-sochat-milroy.pdf)
+- [x] [Container Training Slides](https://drive.google.com/drive/folders/1_mTBBc98TEX3XFpNp0rqoqj1VjN9TKoO)
+- [ ] [Containers as Jupyter Kernels](https://docs.nersc.gov/services/jupyter/how-to-guides/#how-to-use-a-container-to-run-a-jupyter-kernel)
+- [ ] [Using SPIN to Run Persistent Containers](https://docs.nersc.gov/services/spin/)
+- [ ] [Using uv to package lammps and flux into pip install???](https://sgoel.dev/posts/building-cython-or-c-extensions-using-uv/)

matensemble-0.2.1/src/matensemble/__init__.py

@@ -0,0 +1,15 @@
+"""
+MatEnsemble
+
+MatEnsemble is a Python workflow library for building and running
+high-throughput and dependency-aware workflows on HPC systems.
+It lets users define delayed Python and executable chores, connect them
+through data dependencies, and submit them for execution with Flux.
+"""
+
+__author__ = ["Soumendu Bagchi", "Kaleb Duchesneau"]
+__package__ = "matensemble"
+
+# Re-export core data model types at the package root for convenience and for
+# backwards compatibility with code/tests that import from `matensemble`.
+from .model import OutputReference, Resources, ChoreType  # noqa: E402,F401

matensemble-0.2.1/src/matensemble/chore.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+import json
+
+import networkx as nx
+import shlex
+
+from pathlib import Path
+
+from matensemble.model import ChoreType, Resources
+from matensemble.utils import _json_safe
+
+
+class Chore:
+    """
+    A :obj:`Chore` is what MatEnsemble is built around. :obj:`Job`'s can have two
+    different types. ``PYTHON`` or ``EXECUTABLE``
+
+    Python chores are delayed function calls that will be submitted to the
+    runtime-worker when the :obj:`Chore`'s dependencies are resolved and they are
+    schduled in the queue.
+
+    Executable chores are simply commands that will usually call an Executable script
+    when the chore is scheduled.
+
+    """
+
+    def __init__(
+        self,
+        id: str,
+        command: str | list[str],
+        chore_type: ChoreType,
+        resources: Resources,
+        workdir: Path,
+        func_module: str | None = None,
+        func_qualname: str | None = None,
+        serialized_callable: bytes | None = None,
+        deps: tuple[str, ...] = (),
+        args: tuple = (),
+        kwargs: dict | None = None,
+    ) -> None:
+        """
+        The constructor for a :obj:`Chore`
+
+        Parameters
+        ----------
+        id : str
+            The ID for the :obj:`Chore`
+        command : str, list[str]
+            The command that will be run when the :obj:`Chore` is submitted
+        chore_type: ChoreType
+            Either PYTHON or EXECUTABLE
+        resources : Resources
+            An instance of :obj:`Resources` that holds all the information about
+            what resources are needed to run the :obj:`Chore`
+        workdir : Path
+            The Path to the directory where the output of the :obj:`Chore` will be
+            handled
+        func_module : str
+            The module where the function definition is if the type of the :obj:`Chore`
+            is PYTHON
+        func_qualname : str
+            The name of the function if the type of the :obj:`Chore` is PYTHON
+        serialized_callable : bytes
+            The original function that was wrapped stored as bytes
+        deps : tuple[str, ...]
+            A tupele of chore-id's which results this :obj:`Chore
+        args : tuple
+            The arguments to give the function if type is PYTHON
+        kwargs : dict
+            The key-word arguments to give the function if flaovr is PYTHON
+        """
+
+        self.id = id
+        self.command = (
+            shlex.split(command) if isinstance(command, str) else list(command)
+        )
+
+        if chore_type == ChoreType.PYTHON:
+            if serialized_callable is None and not (func_module and func_qualname):
+                raise ValueError(
+                    "Python chores require either serialized_callable or func_module+func_qualname"
+                )
+
+        self.chore_type = chore_type
+        self.resources = resources
+        self.workdir = workdir.resolve()
+        self.spec_path = self.workdir / "chore.pkl"
+
+        self.func_module = func_module
+        self.func_qualname = func_qualname
+        self.serialized_callable = serialized_callable
+        
+        self.deps = deps
+        self.args = args
+        self.kwargs = {} if kwargs is None else kwargs
+
+    def graph(self) -> nx.DiGraph:
+        return nx.DiGraph()
+
+    def _to_debug_dict(self) -> dict:
+        return {
+            "id": self.id,
+            "command": self.command,
+            "chore_type": _json_safe(self.chore_type),
+            "resources": {
+                "num_tasks": self.resources.num_tasks,
+                "cores_per_task": self.resources.cores_per_task,
+                "gpus_per_task": self.resources.gpus_per_task,
+                "mpi": self.resources.mpi,
+                "env": _json_safe(self.resources.env),
+                "inherit_env": self.resources.inherit_env,
+            },
+            "spec_file": str(self.spec_path),
+            "func_module": self.func_module,
+            "func_qualname": self.func_qualname,
+            "has_serialized_callable": self.serialized_callable is not None,
+            "deps": list(self.deps),
+            "args": _json_safe(self.args),
+            "kwargs": _json_safe(self.kwargs),
+        }
+
+    def _write_debug_json(self) -> None:
+        """
+        The :obj:`Chore` is pickled at runtime to be used later on, but it is also
+        written as json for debugging.
+        """
+
+        debug_file = self.spec_path.parent / "chore.json"
+        debug_file.parent.mkdir(parents=True, exist_ok=True)
+        with debug_file.open("w") as f:
+            json.dump(self._to_debug_dict(), f, indent=2)
+
+    def __str__(self) -> str:
+        """
+        Return the :obj:`Chore` as a JSON string.
+        """
+
+        return json.dumps(self._to_debug_dict(), indent=2)