mamba

Mamba provides a powerful environment management system that includes both the conda and mamba package managers. While conda is widely known for managing environments and packages in the Python ecosystem, mamba offers a faster, more efficient alternative for solving package dependencies and installations. With mamba, users can benefit from significant performance improvements, especially when managing large environments, while still maintaining compatibility with conda's extensive ecosystem. Both tools can be used interchangeably within the Mamba framework, offering flexibility for users depending on their needs.

Mamba Package Manager and Conda Distribution

To use Mamba, you need to load the appropriate module for the Python distribution version you want to use. The following example shows how to load the Mamba module (the $ is the prompt; do not type it):

$ module load mamba

To check which versions are available, use the module command:

$ module available mamba

Then, load the version you need. Please see our page on Lmod for more information on loading modules.

Installing Packages with Mamba

Mamba is an optimized package manager, a drop-in replacement for conda. To create a new environment and install packages:

$ mamba create -n myenv python=3.10

After creating an environment, activate it:

$ source activate myenv

You can then install additional packages:

$ mamba install numpy scipy

To deactivate the environment:

$ source deactivate

Installing Packages with Conda

To create a new environment and install packages:

$ conda create -n myenv python=3.10

After creating an environment, activate it:

$ source activate myenv

You can then install additional packages:

$ conda install numpy scipy

To deactivate the environment:

$ source deactivate

Running Python Scripts

To run Python at a prompt, simply type:

$ python

You may also run Python scripts directly from the command line. For example, if you have a script my_script.py in the current directory:

$ python ./my_script.py

As with MATLAB, you can run Python scripts on a login node for smaller, less computationally intensive tasks. Ensure your program does not exceed memory or CPU limits (as a guideline, no more than 4 GB of memory and a runtime of less than 5–10 minutes). For large, long-running computations, submit the job through the batch system.

Submitting Jobs with Mamba

When running Python scripts on the cluster, especially for larger jobs, it’s recommended to submit them through SLURM or other job scheduling systems. Suppose you have a Python program my_script.py. You can run this in batch mode by creating a SLURM script that looks like this:

#!/bin/bash
#SBATCH --job-name=my_job
#SBATCH --output=output.txt
#SBATCH --ntasks=1
#SBATCH --time=01:00:00
#SBATCH --mem=4G

# Load the Mamba module and activate environment
module load mamba
source activate myenv

# Run the Python script
python my_script.py

Submit this script with the sbatch command:

$ sbatch my_job.slurm

Parallel Computing with Python and Mamba

Python supports parallel computing using several libraries like multiprocessing, Dask, and mpi4py. These libraries allow Python code to take advantage of multiple CPUs available on a node or across nodes in a cluster.

Using Multiprocessing

For simple parallelism using multiple cores on a single machine, Python’s built-in multiprocessing library can be used. For example, in your script:

import multiprocessing as mp

def worker(task):
    # Your task here
    pass

if __name__ == '__main__':
    pool = mp.Pool(mp.cpu_count())  # Utilize all available CPUs
    results = pool.map(worker, tasks)
    pool.close()
    pool.join()

Distributed Computing with Dask

For distributed computing across multiple nodes or larger datasets, Dask is a powerful framework. Dask enables parallelism for many Python libraries like numpy, pandas, and more.

Install Dask using Mamba:

$ mamba install dask distributed

You can submit Dask jobs via SLURM. Here’s a basic SLURM submission script for a Dask job:

#!/bin/bash
#SBATCH --job-name=dask_job
#SBATCH --output=dask_output.txt
#SBATCH --ntasks=4
#SBATCH --time=02:00:00
#SBATCH --mem=16G

module load mamba
source activate myenv

# Launch Dask Scheduler and Workers
dask-scheduler --scheduler-file scheduler.json &
sleep 5
dask-worker --scheduler-file scheduler.json --nthreads 4 &
dask-worker --scheduler-file scheduler.json --nthreads 4 &

You can then connect to this Dask cluster from your Python code using:

from dask.distributed import Client
client = Client(scheduler_file='scheduler.json')

# Submit tasks to the cluster

Using mpi4py for MPI

For more advanced distributed computing across multiple nodes, you can use mpi4py, a Python binding for MPI. Install it via Mamba:

$ mamba install mpi4py

In your Python script, use MPI as follows:

from mpi4py import MPI

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

print(f'Hello from rank {rank} out of {size}')

A SLURM script for running an MPI Python program:

#!/bin/bash
#SBATCH --job-name=mpi_python
#SBATCH --ntasks=4
#SBATCH --time=01:00:00
#SBATCH --mem=8G

module load mamba
source activate myenv

# Run the MPI program
mpirun python mpi_script.py

Submit this script with:

$ sbatch mpi_job.slurm

Cleaning Up

After your parallel job has completed, ensure all resources are released. For instance, when using Dask, shut down your workers and scheduler at the end of your job, or use client.close() in your Python script.

Pre-installed packages