feat: Implement resource requirements for components

.github/workflows/lint-test.yaml

@@ -75,7 +75,7 @@ jobs:
   test-unit:
     name: Tests - unit
     runs-on: ubuntu-latest
-    timeout-minutes: 5
+    timeout-minutes: 8
     strategy:
       matrix:
         python_version: [3.12, 3.13]
@@ -113,7 +113,7 @@ jobs:
   test-integration:
     name: Tests - integration
     runs-on: ubuntu-latest
-    timeout-minutes: 6
+    timeout-minutes: 10
     strategy:
       matrix:
         python_version: [3.12, 3.13]
@@ -157,7 +157,7 @@ jobs:
   test-integration-tuner:
     name: Tests - integration:tuner
     runs-on: ubuntu-latest
-    timeout-minutes: 5
+    timeout-minutes: 8
     strategy:
       matrix:
         python_version: [3.12, 3.13]

docs/api/schemas/schemas.md

@@ -1 +1 @@
-::: plugboard.schemas
+::: plugboard_schemas

docs/examples/tutorials/running-in-parallel.md

@@ -70,3 +70,68 @@ Specifying the process type and channel builder type in the YAML is the only cha
 
 1. Tell Plugboard to use a [`RayProcess`][plugboard.process.RayProcess] instead of the default [`LocalProcess`][plugboard.process.LocalProcess].
 2. Also change the connector builder to [`RayConnector`][plugboard.connector.RayConnector], which will build [`RayChannel`][plugboard.connector.RayChannel] objects when creating the `Process`.
+
+## Specifying resource requirements
+
+When running components on Ray, you can specify resource requirements for each component to control how Ray allocates computational resources. This is particularly useful when you have components with different resource needs (e.g., CPU-intensive vs GPU-intensive tasks) and you are running on a Ray cluster.
+
+!!! tip
+    Normally Ray will start automatically when you are using Plugboard locally. If you want to start a separate Ray instance, for example so that you can control the configuration options, you can launch it from the [CLI](https://docs.ray.io/en/latest/ray-core/starting-ray.html). For example, this command will start a Ray instance with enough resources to run the example below.
+
+    ```sh
+    uv run ray start --head --num-cpus=4 --num-gpus=1 --resources='{"custom_hardware": 5}'
+    ```
+
+### Declaring resources at component definition
+
+The recommended way to specify resource requirements is to declare them as a class attribute when defining your component. This makes the resource requirements explicit and part of the component's definition:
+
+```python
+from plugboard.component import Component
+from plugboard.schemas import Resource
+
+class CPUIntensiveTask(Component):
+    """Component that requires more CPU resources."""
+
+    io = IO(inputs=["x"], outputs=["y"])
+    resources = Resource(cpu=2.0)  # Declare resources at class level
+
+    async def step(self) -> None:
+        # Your component logic here
+        pass
+```
+
+### Overriding resources at instantiation
+
+You can also override resource requirements when creating component instances. This is useful when you want to use the same component class with different resource requirements:
+
+```python
+# Override the class-level resource requirements
+component = CPUIntensiveTask(
+    name="my-task",
+    resources=Resource(cpu=4.0)  # Override to use 4 CPUs instead of 2
+)
+```
+
+### Example
+
+For example, you can specify [`Resource`][plugboard.schemas.Resource] requirements like this when defining components:
+
+```python
+--8<-- "examples/tutorials/004_using_ray/resources_example.py:resources"
+```
+
+Or override them in YAML configuration:
+
+```yaml
+--8<-- "examples/tutorials/004_using_ray/resources-example.yaml:10:"
+```
+
+1. Override DataProducer to require 1 CPU (instead of the default 0.001).
+2. CPUIntensiveTask already declares cpu: 2.0 at the class level, so this matches the class definition.
+3. Add 512MB memory requirement to CPUIntensiveTask (extending the class-level resources).
+4. Requires 0.5 CPU, specified in Kubernetes-style format (500 milli CPUs). This matches the class-level declaration.
+5. Requires 1 GPU, matching the class-level declaration.
+6. Add a custom resource called `custom_hardware`. This needs to be specified in the configuration of your Ray cluster to make it available.
+
+See the [Ray documentation](https://docs.ray.io/en/latest/ray-core/scheduling/resources.html) for more information about specifying resource requirements.

examples/tutorials/004_using_ray/resources-example.yaml

@@ -0,0 +1,47 @@
+# Example YAML configuration with resource requirements
+#
+# This configuration demonstrates how to specify resource requirements
+# for components in a Plugboard process. Resources can be specified as:
+# - Numerical values: cpu: 2.0
+# - Milli-units: cpu: "250m" (equals 0.25)
+# - Memory units: memory: "100Mi" (equals 100 * 1024 * 1024 bytes)
+#
+# Note: Resources can be declared at the component class level (recommended)
+# or overridden in the YAML configuration as shown below.
+plugboard:
+  process:
+    type: plugboard.process.RayProcess
+    connector_builder:
+      type: plugboard.connector.RayConnector
+    args:
+      name: resource-example-process
+      components:
+        - type: examples.tutorials.004_using_ray.resources_example.DataProducer
+          args:
+            name: producer
+            iters: 10
+            resources:
+              cpu: 1.0   # (1)!
+        - type: examples.tutorials.004_using_ray.resources_example.CPUIntensiveTask
+          args:
+            name: cpu-task
+            # CPUIntensiveTask has class-level resources (cpu: 2.0)
+            # Override to use more memory
+            resources:
+              cpu: 2.0  # (2)!
+              memory: "512Mi"  # (3)!
+        - type: examples.tutorials.004_using_ray.resources_example.GPUTask
+          args:
+            name: gpu-task
+            # GPUTask has class-level resources (cpu: "500m", gpu: 1)
+            # Can override or extend with custom resources
+            resources:
+              cpu: "500m"  # (4)!
+              gpu: 1  # (5)!
+              resources:
+                custom_hardware: 2  # (6)!
+      connectors:
+        - source: producer.output
+          target: cpu-task.x
+        - source: cpu-task.y
+          target: gpu-task.data

examples/tutorials/004_using_ray/resources_example.py

@@ -0,0 +1,95 @@
+"""Example demonstrating resource requirements for components in RayProcess."""
+
+import asyncio
+import typing as _t
+
+import ray
+
+from plugboard.component import Component, IOController as IO
+from plugboard.connector import RayConnector
+from plugboard.process import RayProcess
+from plugboard.schemas import ComponentArgsDict, ConnectorSpec, Resource
+
+
+class CPUIntensiveTask(Component):
+    """Component that requires more CPU resources.
+
+    Resource requirements are declared as a class attribute.
+    """
+
+    io = IO(inputs=["x"], outputs=["y"])
+    resources = Resource(cpu=2.0)  # Declare resources at class level
+
+    async def step(self) -> None:
+        """Execute CPU-intensive computation."""
+        # Simulate CPU-intensive work
+        result = sum(i**2 for i in range(int(self.x * 10000)))
+        self.y = result
+
+
+class GPUTask(Component):
+    """Component that requires GPU resources.
+
+    Resource requirements are declared as a class attribute.
+    """
+
+    io = IO(inputs=["data"], outputs=["result"])
+    resources = Resource(cpu="500m", gpu=1)  # Declare resources at class level
+
+    async def step(self) -> None:
+        """Execute GPU computation."""
+        # Simulate GPU computation
+        self.result = self.data * 2
+
+
+class DataProducer(Component):
+    """Produces data for processing."""
+
+    io = IO(outputs=["output"])
+
+    def __init__(self, iters: int, **kwargs: _t.Unpack[ComponentArgsDict]) -> None:
+        """Initialize DataProducer with iteration count."""
+        super().__init__(**kwargs)
+        self._iters = iters
+
+    async def init(self) -> None:
+        """Initialize the data sequence."""
+        self._seq = iter(range(self._iters))
+
+    async def step(self) -> None:
+        """Produce the next data value."""
+        try:
+            self.output = next(self._seq)
+        except StopIteration:
+            await self.io.close()
+
+
+async def main() -> None:
+    """Run the process with resource-constrained components."""
+    # --8<-- [start:resources]
+    # Resources can be declared at the class level (see CPUIntensiveTask and GPUTask above)
+    # or overridden when instantiating components
+    process = RayProcess(
+        components=[
+            DataProducer(name="producer", iters=5),  # Uses default resources
+            CPUIntensiveTask(name="cpu-task"),  # Uses class-level resources (2.0 CPU)
+            GPUTask(name="gpu-task"),  # Uses class-level resources (0.5 CPU, 1 GPU)
+        ],
+        connectors=[
+            RayConnector(spec=ConnectorSpec(source="producer.output", target="cpu-task.x")),
+            RayConnector(spec=ConnectorSpec(source="cpu-task.y", target="gpu-task.data")),
+        ],
+    )
+    # --8<-- [end:resources]
+
+    async with process:
+        await process.run()
+
+    print("Process completed successfully!")
+
+
+if __name__ == "__main__":
+    if not ray.is_initialized():
+        # Ray must be initialised with the necessary resources
+        ray.init(num_cpus=5, num_gpus=1, resources={"custom_hardware": 10}, include_dashboard=True)
+    asyncio.run(main())

mkdocs.yaml

@@ -33,7 +33,7 @@ plugins:
     default_handler: python
     handlers:
       python:
-        paths: [src]
+        paths: [plugboard, plugboard-schemas]
         options:
           docstring_style: google
           show_source: false
@@ -113,6 +113,7 @@ watch:
 - docs
 - examples
 - plugboard
+- plugboard-schemas
 - README.md
 - CONTRIBUTING.md
 

plugboard-schemas/plugboard_schemas/__init__.py

@@ -9,7 +9,7 @@
 from importlib.metadata import version
 
 from ._common import PlugboardBaseModel
-from .component import ComponentArgsDict, ComponentArgsSpec, ComponentSpec
+from .component import ComponentArgsDict, ComponentArgsSpec, ComponentSpec, Resource
 from .config import ConfigSpec, ProcessConfigSpec
 from .connector import (
     DEFAULT_CONNECTOR_CLS_PATH,
@@ -77,6 +77,7 @@
     "ProcessArgsDict",
     "ProcessArgsSpec",
     "RAY_STATE_BACKEND_CLS_PATH",
+    "Resource",
     "StateBackendSpec",
     "StateBackendArgsDict",
     "StateBackendArgsSpec",