Skip to content

Commit 1d0f046

Browse files
drubinsteindrubinstein
committed
Add mamba support
1 parent f7988ec commit 1d0f046

File tree

11 files changed

+142
-21
lines changed

11 files changed

+142
-21
lines changed

‎.github/workflows/mamba.yml‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
name: Test conda
2+
on: [push, pull_request]
3+
4+
jobs:
5+
test:
6+
name: test ${{ matrix.py }} - ${{ matrix.os }}
7+
runs-on: ${{ matrix.os }}-latest
8+
strategy:
9+
fail-fast: false
10+
matrix:
11+
os:
12+
- Ubuntu
13+
py:
14+
- "3.8"
15+
16+
steps:
17+
- name: Setup python for test ${{ matrix.py }}
18+
uses: actions/setup-python@v3
19+
with:
20+
python-version: ${{ matrix.py }}
21+
- uses: conda-incubator/setup-miniconda@v2
22+
with:
23+
auto-update-conda: true
24+
python-version: ${{ matrix.python-version}}
25+
- name: Install mamba
26+
run: conda install mamba -n base -c conda-forge -y
27+
- uses: actions/checkout@v3
28+
- name: Upgrade pip
29+
run: python -m pip install -U pip
30+
- name: Install dependencies
31+
run: python -m pip install -e '.[test]'
32+
- name: Run pytest
33+
run: pytest

‎README.md‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,14 +11,15 @@ Python Compose is a command line tool and library for spinning up many long-runn
1111

1212
Currently, Python Compose supports the following environments:
1313
- [conda](https://docs.conda.io/en/latest/)
14+
- [mamba](https://mamba.readthedocs.io/en/latest/index.html)
1415
- [pyenv-virtualenv](https://github.com/pyenv/pyenv-virtualenv)
1516
- [venv](https://docs.python.org/3/library/venv.html)
1617

1718
In the future, we wish to support:
1819

1920
- [poetry](https://python-poetry.org/)
2021
- [pipenv](https://pipenv.pypa.io/en/latest/)
21-
- [mamba](https://mamba.readthedocs.io/en/latest/index.html)
22+
- [pyenv](https://github.com/pyenv/pyenv)
2223

2324
## Installation
2425

‎examples/mamba/__init__.py‎

Whitespace-only changes.

‎examples/mamba/example.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
from python_compose import compose
2+
from python_compose.unit.conda import MambaUnit
3+
4+
compose.compose(
5+
[
6+
MambaUnit(
7+
name=f"httpd_{i}", requirements=[], command=["python3", "httpd.py", str(8080 + i)]
8+
)
9+
for i in range(3)
10+
]
11+
)

‎examples/mamba/httpd.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
import argparse
2+
from http.server import BaseHTTPRequestHandler, HTTPServer
3+
4+
if __name__ == "__main__":
5+
parser = argparse.ArgumentParser()
6+
parser.add_argument("port", type=int)
7+
args = parser.parse_args()
8+
9+
print(args)
10+
server_address = ("", args.port)
11+
httpd = HTTPServer(server_address, BaseHTTPRequestHandler)
12+
httpd.serve_forever()

‎pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ dependencies = [
2929
"pydantic-yaml",
3030
"ruamel.yaml",
3131
"typing-extensions",
32+
"types-PyYAML"
3233
]
3334

3435
[project.urls]

‎python_compose/compose.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55

66
from python_compose.spec import Unit, YamlSpec
77
from python_compose.unit.compose_unit import ComposeUnit
8-
from python_compose.unit.conda import CondaUnit
8+
from python_compose.unit.conda import CondaUnit, MambaUnit
99
from python_compose.unit.pyenv_virtualenv import PyenvVirtualenvUnit
1010
from python_compose.unit.venv import VenvUnit
1111

@@ -66,6 +66,8 @@ def pydantic_to_units(units: List[Unit]) -> List[ComposeUnit]:
6666
kwargs = {k: v for k, v in unit.dict().items() if k != "unit_type"}
6767
if unit.unit_type == "conda":
6868
ret.append(CondaUnit(**kwargs))
69+
elif unit.unit_type == "mamba":
70+
ret.append(MambaUnit(**kwargs))
6971
elif unit.unit_type == "pyenv-virtualenv":
7072
ret.append(PyenvVirtualenvUnit(**kwargs))
7173
elif unit.unit_type == "venv":

‎python_compose/spec.py‎

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,9 @@
88

99
# TODO: Add classmethods to convert these to their a
1010
# TODO: Move models to the associated unit class
11-
class CondaUnitModel(BaseModel):
12-
"""The definition for running a Conda Unit."""
11+
class AnacondaUnitModelMixin(BaseModel):
12+
"""The definition for running an Anaconda-compatible Unit."""
1313

14-
unit_type: Literal["conda"]
15-
"""Definition that this is a conda model."""
1614
name: str
1715
"""The name of the environment."""
1816
requirements: Union[pathlib.Path, List[str]] = []
@@ -23,6 +21,20 @@ class CondaUnitModel(BaseModel):
2321
"""The optional working directory for the script being run."""
2422

2523

24+
class CondaUnitModel(AnacondaUnitModelMixin):
25+
"""The definition for running a conda Unit."""
26+
27+
unit_type: Literal["conda"]
28+
"""Definition that this is a conda model."""
29+
30+
31+
class MambaUnitModel(AnacondaUnitModelMixin):
32+
"""The definition for running a mamba Unit."""
33+
34+
unit_type: Literal["mamba"]
35+
"""Definition that this is a mamba model."""
36+
37+
2638
class PyenvVirtualenvUnitModel(BaseModel):
2739
"""The definition for running a pyenv Unit."""
2840

@@ -58,7 +70,7 @@ class VenvUnitModel(BaseModel):
5870

5971

6072
Unit = Annotated[
61-
Union[CondaUnitModel, PyenvVirtualenvUnitModel, VenvUnitModel],
73+
Union[CondaUnitModel, MambaUnitModel, PyenvVirtualenvUnitModel, VenvUnitModel],
6274
Field(discriminator="unit_type"),
6375
]
6476
"""The collection of Unit models that we will be able to deserialize."""

‎python_compose/unit/conda.py‎

Lines changed: 29 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,11 @@
66
from python_compose.unit.compose_unit import ComposeUnit
77

88

9-
class CondaUnit(ComposeUnit):
10-
"""A Compose Unit for creating Conda environments and running scripts in them."""
9+
class AnacondaUnit(ComposeUnit):
10+
"""A base Compose Unit for creating anaconda-compatible environments and
11+
running scripts in them."""
12+
13+
EXECUTABLE_NAME = ""
1114

1215
def __init__(
1316
self,
@@ -19,10 +22,10 @@ def __init__(
1922
"""
2023
2124
Args:
22-
name: Name of the conda environment to create.
25+
name: Name of the environment to create.
2326
requirements: Either a path to a requirements.txt file or a list of requirements to
2427
install.
25-
command: The command to run in the conda environment
28+
command: The command to run in the environment
2629
working_dir: The optional working directory for the command being run.
2730
"""
2831
self.name = name
@@ -34,33 +37,49 @@ def create(self) -> None:
3437
"""Function for creating a virtual environment."""
3538
envs = [
3639
row.split()[0]
37-
for row in subprocess.check_output(["conda", "env", "list"]).decode().split("\n")[2:]
40+
for row in subprocess.check_output([self.EXECUTABLE_NAME, "env", "list"])
41+
.decode()
42+
.split("\n")[2:]
3843
if row
3944
]
4045
if self.name in envs:
4146
warnings.warn(f"Skipping pyenv venv creation for {self.name}. Venv already exists.")
4247
else:
43-
subprocess.check_call(["conda", "create", "-n", self.name, "-y"])
48+
subprocess.check_call([self.EXECUTABLE_NAME, "create", "-n", self.name, "-y"])
4449

4550
def install_requirements(self) -> None:
4651
"""Function to install any and all requirements for running a script in the virtual
4752
environment."""
4853
if isinstance(self.requirements, list) and self.requirements:
49-
subprocess.check_call(["conda", "install", "-n", self.name] + self.requirements)
54+
subprocess.check_call(
55+
[self.EXECUTABLE_NAME, "install", "-n", self.name] + self.requirements
56+
)
5057
elif isinstance(self.requirements, pathlib.Path):
5158
subprocess.check_call(
52-
["conda", "install", "-n", self.name, "-f", str(self.requirements)]
59+
[self.EXECUTABLE_NAME, "install", "-n", self.name, "-f", str(self.requirements)]
5360
)
5461

5562
def start(self) -> None:
5663
"""Function to start a script in the previously created virtual environment."""
5764
p = subprocess.Popen(
58-
["conda", "run", "-n", self.name, "--no-capture-output"]
65+
[self.EXECUTABLE_NAME, "run", "-n", self.name, "--no-capture-output"]
5966
+ (["--cwd", str(self.working_dir)] if self.working_dir else [])
6067
+ self.command
6168
)
6269
p.communicate()
6370

6471
def clean(self) -> None:
6572
"""Function to erase any pre-existing files to be recreated by a Python Compose Unit."""
66-
subprocess.check_call(["conda", "remove", "-n", self.name, "--all"])
73+
subprocess.check_call([self.EXECUTABLE_NAME, "remove", "-n", self.name, "--all"])
74+
75+
76+
class CondaUnit(AnacondaUnit):
77+
"""A Compose Unit for creating conda environments and running scripts in them."""
78+
79+
EXECUTABLE_NAME = "conda"
80+
81+
82+
class MambaUnit(AnacondaUnit):
83+
"""A Compose Unit for creating mamba environments and running scripts in them."""
84+
85+
EXECUTABLE_NAME = "mamba"

‎tests/config.yaml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,3 +12,6 @@ units:
1212
- unit_type: "conda"
1313
name: httpd_2
1414
command: ["python3", "../conda/httpd.py", "8082"]
15+
- unit_type: "mamba"
16+
name: httpd_3
17+
command: ["python3", "../conda/httpd.py", "8082"]

0 commit comments

Comments
 (0)