RayClusterResource Dagster resource implementation that provision/deprovision a new Ray cluster on demand

ray_cluster_resource_dagster.py

from typing import Optional

from dagster import AssetExecutionContext, ConfigurableResource, Definitions, InitResourceContext, asset

from ray.autoscaler.sdk import create_or_update_cluster, run_on_cluster, teardown_cluster


class RayClusterResource(ConfigurableResource):
    cluster_config: str

    def setup_for_execution(self, context: InitResourceContext) -> None:
        # https://docs.ray.io/en/latest/_modules/ray/autoscaler/sdk/sdk.html#create_or_update_cluster
        # TODO suffix ray cluster name with run_id
        # context.dagster_run
        context.log_manager.info("Starting new Ray cluster")
        create_or_update_cluster(cluster_config=self.cluster_config)

    def teardown_after_execution(self, context: InitResourceContext) -> None:
        # https://docs.ray.io/en/latest/_modules/ray/autoscaler/sdk/sdk.html#teardown_cluster
        teardown_cluster(cluster_config=self.cluster_config)

    def run(self, cmd: str) -> Optional[str]:
        # https://docs.ray.io/en/latest/_modules/ray/autoscaler/sdk/sdk.html#run_on_cluster
        return run_on_cluster(cluster_config=self.cluster_config, cmd=cmd)


@asset
def webdataset(ray_cluster: RayClusterResource, context: AssetExecutionContext) -> None:
    # recorded metadata can be customized
    # metadata = {
    #     "num_records": len(df),
    #     "preview": MetadataValue.md(df[["title", "by", "url"]].to_markdown()),
    # }
    # context.add_output_metadata(metadata=metadata)
    ray_cluster.run(cmd="""python3.11 -c 'print("Hello World")'""")


defs = Definitions(
    assets=[webdataset],
    resources={
        "ray_cluster": RayClusterResource(cluster_config="/etc/ray/example_ray_config.yml"),
    },
)