Skip to content

scienting/raygent

Repository files navigation

raygent

Parallelism, Delegated

Raygent simplifies parallel execution in Python by providing an intuitive interface to Ray's distributed computing framework. It removes boilerplate code and offers a modular approach to managing parallel tasks with directed acyclic graphs (DAG), making it easier to scale your workflow across multiple computational cores and nodes.

Caution

Raygent is under active development. Breaking changes are could happen without warning.

Features

  • Simple Task-Based API: Define your computational logic once and automatically scale to any size.
  • Flexible Execution Modes: Run tasks either in parallel using Ray or sequentially with a single parameter.
  • Resource Optimization: Automatically detect available CPU cores and manage resource allocation.
  • Batched Processing: Efficiently process data in optimal batch sizes.
  • DAG workflows: Easy-to-use workflows using parallelized computational tasks (nodes) and data transformations (edges).

Quick Start

Defining Tasks

All workflows are built from raygent Tasks that specify an independent calculation. At minimum, Tasks must have a do method that performs the desired computation. Tasks are not limited to native Python; they can use NumPy arrays, Polars DataFrames, anything.

from typing import override

from raygent import Task

class SquareTask(Task):
    @override
    def do(self, batch: list[int]) ->  list[int]:
        return [item ** 2 for item in batch]

class PrefactorTask(Task):
    @override
    def do(self, batch: list[int], factor: int = 1) -> list[int]:
        return [factor * x for x in batch]

class CombineTask(Task):
    @override
    def do(self, a: list[int], b: list[int]) -> dict[str, list[int]]:
        return {"a": a, "b": b}

class SumTask(Task):
    @override
    def do(self, batch: dict[str, list[int]]) ->  list[int]:
        return [a + b for a, b in zip(batch["a"], batch["b"])]

Tasks must follow a few rules:

  • Positional arguments take batches of data, not a single element.
  • Only data can be passed through positional arguments; all other parameters must be specified with keyword arguments.
  • Each Task should represent an isolated computation (i.e., no Task can call another Task).
  • Internal multithreading in do methods should be explicitly annotated in the docstring.

Abiding by these rules will make creating workflows with a DAG a breeze.

Important

DAG workflows do not restrict Tasks.do() to a single thread; Tasks could perform their own multithreading. This is a limitation of ray, so take this into account when specifying resources for DAG nodes.

Creating workflows

Here is an example directed acyclic graph (DAG) workflow. It creates two source queues (source_1 and source_2) that inject data using our Tasks defined in the previous section.

flowchart LR
    A[source_1] --> B((PrefactorTask))
    B --> C((SquareTask))
    D[source_2] --> E((SquareTask))

    C --> F((CombineTask))
    E --> F

    F --> G((SumTask))

    G --> I([sink_1])
    E --> H([sink_2])
Loading

This DAG also provides two sinks to receive messages from: sink_1 is our final processed data, and sink_2 provides messages from intermediate nodes. Do note that queues consume computational resources; adding an excessive number of sinks will reduce your workflow's parallelization.

Below is how we build this DAG in raygent.

from raygent.workflow import DAG

dag = DAG()

# Add source nodes we can send data into the DAG
# Returns a source node and a queue
source_n1, source_1 = dag.add_source()
source_n2, source_2 = dag.add_source()

# Add fully-connected nodes to process our workflow
#   Top
prefactor_n = dag.add(PrefactorTask(), inputs=source_n1, task_kwargs={"factor": 2})
square_n1 = dag.add(SquareTask(), inputs=prefactor_n)
#   Bottom
square_n2 = dag.add(SquareTask(), inputs=source_n2)
#   Merged
comb = dag.add(CombineTask(), inputs=(square_n1, square_n2))
summed = dag.add(SumTask(), inputs=comb)

# Attach sinks to get data out of our DAG
sink_1 = dag.add_sink(summed)
sink_2 = dag.add_sink(square_n2)

Note

All inputs to a node must be defined before adding the task to the DAG. Defining nodes from source to sink helps prevent cycles.

Running workflows

Injecting data into sources must be done with DAG.stream() to process data in batches. It will handle batching into source queues and returning messages from all sinks.

data1 = [1, 2, 3, 4]  # For source_1
data2 = [5, 6, 7, 8]  # For source_2

# Workflow starts waiting for source messages
dag.start()

# Streams data into sources and collects any sink messages
# Order of data specifies which source queue it goes into
for q_idx, msg in dag.stream(
    data1,
    data2,
    source_queues=(source_1, source_2),
    sink_queues=(sink_1, sink_2),
    batch_size=2,
    max_inflight=100,
    sink_wait=0.01,
):
    print(f"from sink #{q_idx} -> batch_idx={msg.index}, payload={msg.payload}")

# Releases all resources by terminating all DAG nodes
dag.stop()

# sink_1 payloads: [29, 52], [85, 128]
# sink_2 payloads: [25, 36], [49, 64]

Tip

Sinks operate using a first-in, first-out (FIFO) queue. Batches are processed when finished, not necessarily in the order they are sent into the DAG.

Installation

You can install raygent directly from the GitHub repository. First, clone the repository.

git clone git@github.com:scienting/raygent.git

Install raygent using pip after moving into the directory.

pip install .

This will install all dependencies and raygent into your current Python environment. To enable raygent's full parallelization capabilities with DAG workflows, install it with the workflow extra.

pip install .[workflow]

Development

We use Pixi to manage environments and dependencies. After installing Pixi, clone the repo and run:

pixi install
pixi shell -e dev

This gives you a fully configured environment with testing, linting, formatting, and build tools. See the development guide for the complete walkthrough, including how to run tests, build the package, and publish releases.

License

raygent is licensed under the Prosperity Public License 3.0.0.

Noncommercial use is free. Academic researchers, university labs, nonprofits, and individual learners can use, modify, and distribute this software and any datasets it generates at no cost. If you use raygent in published research, a citation is appreciated.

Commercial use requires a paid license. The Prosperity Public License includes a 30-day trial period for commercial evaluation. After that, for-profit companies using raygent or its outputs in commercial products, services, or internal operations need a commercial license. This includes using the generated datasets in proprietary ML pipelines, commercial drug discovery workflows, or products built on the cleaned data.

Revenue from commercial licenses funds continued development, data validation, and maintenance of this project. To purchase a commercial license, contact us@scient.ing.

About

Parallelism, Delegated

Topics

Resources

License

Stars

0 stars

Watchers

1 watching

Forks

Releases

No releases published

Contributors

Languages