Repository: PrefectHQ/prefect

Issue:
Migrate `prefect-ray` to core



Discussion:

Resolving patch:
diff --git a/src/integrations/prefect-ray/prefect_ray/__init__.py b/src/integrations/prefect-ray/prefect_ray/__init__.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-ray/prefect_ray/__init__.py
@@ -0,0 +1,6 @@
+from . import _version
+
+__version__ = _version.__version__
+
+
+from .task_runners import RayTaskRunner  # noqa
diff --git a/src/integrations/prefect-ray/prefect_ray/context.py b/src/integrations/prefect-ray/prefect_ray/context.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-ray/prefect_ray/context.py
@@ -0,0 +1,71 @@
+"""
+Contexts to manage Ray clusters and tasks.
+"""
+
+from contextlib import contextmanager
+from typing import Any, Dict
+
+from pydantic import VERSION as PYDANTIC_VERSION
+
+from prefect.context import ContextModel, ContextVar
+
+if PYDANTIC_VERSION.startswith("2."):
+    from pydantic.v1 import Field
+else:
+    from pydantic import Field
+
+
+class RemoteOptionsContext(ContextModel):
+    """
+    The context for Ray remote_options management.
+
+    Attributes:
+        current_remote_options: A set of current remote_options in the context.
+    """
+
+    current_remote_options: Dict[str, Any] = Field(default_factory=dict)
+
+    @classmethod
+    def get(cls) -> "RemoteOptionsContext":
+        """
+        Return an empty `RemoteOptionsContext`
+        instead of `None` if no context exists.
+        """
+        return cls.__var__.get(RemoteOptionsContext())
+
+    __var__ = ContextVar("remote_options")
+
+
+@contextmanager
+def remote_options(**new_remote_options: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Context manager to add keyword arguments to Ray `@remote` calls
+    for task runs. If contexts are nested, new options are merged with options
+    in the outer context. If a key is present in both, the new option will be used.
+
+    Yields:
+        The current set of remote options.
+
+    Examples:
+        Use 4 CPUs and 2 GPUs for the `process` task:
+        ```python
+        from prefect import flow, task
+        from prefect_ray.task_runners import RayTaskRunner
+        from prefect_ray.context import remote_options
+
+        @task
+        def process(x):
+            return x + 1
+
+        @flow(task_runner=RayTaskRunner())
+        def my_flow():
+            # equivalent to setting @ray.remote(num_cpus=4, num_gpus=2)
+            with remote_options(num_cpus=4, num_gpus=2):
+                process.submit(42)
+        ```
+    """
+    current_remote_options = RemoteOptionsContext.get().current_remote_options
+    with RemoteOptionsContext(
+        current_remote_options={**current_remote_options, **new_remote_options}
+    ):
+        yield
diff --git a/src/integrations/prefect-ray/prefect_ray/task_runners.py b/src/integrations/prefect-ray/prefect_ray/task_runners.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-ray/prefect_ray/task_runners.py
@@ -0,0 +1,290 @@
+"""
+Interface and implementations of the Ray Task Runner.
+[Task Runners](https://docs.prefect.io/api-ref/prefect/task-runners/)
+in Prefect are responsible for managing the execution of Prefect task runs.
+Generally speaking, users are not expected to interact with
+task runners outside of configuring and initializing them for a flow.
+
+Example:
+    ```python
+    import time
+
+    from prefect import flow, task
+
+    @task
+    def shout(number):
+        time.sleep(0.5)
+        print(f"#{number}")
+
+    @flow
+    def count_to(highest_number):
+        for number in range(highest_number):
+            shout.submit(number)
+
+    if __name__ == "__main__":
+        count_to(10)
+
+    # outputs
+    #0
+    #1
+    #2
+    #3
+    #4
+    #5
+    #6
+    #7
+    #8
+    #9
+    ```
+
+    Switching to a `RayTaskRunner`:
+    ```python
+    import time
+
+    from prefect import flow, task
+    from prefect_ray import RayTaskRunner
+
+    @task
+    def shout(number):
+        time.sleep(0.5)
+        print(f"#{number}")
+
+    @flow(task_runner=RayTaskRunner)
+    def count_to(highest_number):
+        for number in range(highest_number):
+            shout.submit(number)
+
+    if __name__ == "__main__":
+        count_to(10)
+
+    # outputs
+    #3
+    #7
+    #2
+    #6
+    #4
+    #0
+    #1
+    #5
+    #8
+    #9
+    ```
+"""
+
+from contextlib import AsyncExitStack
+from typing import Awaitable, Callable, Dict, Optional
+from uuid import UUID
+
+import anyio
+import ray
+from ray.exceptions import RayTaskError
+
+from prefect.futures import PrefectFuture
+from prefect.states import State, exception_to_crashed_state
+from prefect.task_runners import BaseTaskRunner, R, TaskConcurrencyType
+from prefect.utilities.asyncutils import sync_compatible
+from prefect.utilities.collections import visit_collection
+from prefect_ray.context import RemoteOptionsContext
+
+
+class RayTaskRunner(BaseTaskRunner):
+    """
+    A parallel task_runner that submits tasks to `ray`.
+    By default, a temporary Ray cluster is created for the duration of the flow run.
+    Alternatively, if you already have a `ray` instance running, you can provide
+    the connection URL via the `address` kwarg.
+    Args:
+        address (string, optional): Address of a currently running `ray` instance; if
+            one is not provided, a temporary instance will be created.
+        init_kwargs (dict, optional): Additional kwargs to use when calling `ray.init`.
+    Examples:
+        Using a temporary local ray cluster:
+        ```python
+        from prefect import flow
+        from prefect_ray.task_runners import RayTaskRunner
+
+        @flow(task_runner=RayTaskRunner())
+        def my_flow():
+            ...
+        ```
+        Connecting to an existing ray instance:
+        ```python
+        RayTaskRunner(address="ray://192.0.2.255:8786")
+        ```
+    """
+
+    def __init__(
+        self,
+        address: str = None,
+        init_kwargs: dict = None,
+    ):
+        # Store settings
+        self.address = address
+        self.init_kwargs = init_kwargs.copy() if init_kwargs else {}
+
+        self.init_kwargs.setdefault("namespace", "prefect")
+
+        # Runtime attributes
+        self._ray_refs: Dict[str, "ray.ObjectRef"] = {}
+
+        super().__init__()
+
+    def duplicate(self):
+        """
+        Return a new instance of with the same settings as this one.
+        """
+        return type(self)(address=self.address, init_kwargs=self.init_kwargs)
+
+    def __eq__(self, other: object) -> bool:
+        """
+        Check if an instance has the same settings as this task runner.
+        """
+        if type(self) == type(other):
+            return (
+                self.address == other.address and self.init_kwargs == other.init_kwargs
+            )
+        else:
+            return NotImplemented
+
+    @property
+    def concurrency_type(self) -> TaskConcurrencyType:
+        return TaskConcurrencyType.PARALLEL
+
+    async def submit(
+        self,
+        key: UUID,
+        call: Callable[..., Awaitable[State[R]]],
+    ) -> None:
+        if not self._started:
+            raise RuntimeError(
+                "The task runner must be started before submitting work."
+            )
+
+        call_kwargs, upstream_ray_obj_refs = self._exchange_prefect_for_ray_futures(
+            call.keywords
+        )
+
+        remote_options = RemoteOptionsContext.get().current_remote_options
+        # Ray does not support the submission of async functions and we must create a
+        # sync entrypoint
+        if remote_options:
+            ray_decorator = ray.remote(**remote_options)
+        else:
+            ray_decorator = ray.remote
+
+        self._ray_refs[key] = (
+            ray_decorator(self._run_prefect_task)
+            .options(name=call.keywords["task_run"].name)
+            .remote(sync_compatible(call.func), *upstream_ray_obj_refs, **call_kwargs)
+        )
+
+    def _exchange_prefect_for_ray_futures(self, kwargs_prefect_futures):
+        """Exchanges Prefect futures for Ray futures."""
+
+        upstream_ray_obj_refs = []
+
+        def exchange_prefect_for_ray_future(expr):
+            """Exchanges Prefect future for Ray future."""
+            if isinstance(expr, PrefectFuture):
+                ray_future = self._ray_refs.get(expr.key)
+                if ray_future is not None:
+                    upstream_ray_obj_refs.append(ray_future)
+                    return ray_future
+            return expr
+
+        kwargs_ray_futures = visit_collection(
+            kwargs_prefect_futures,
+            visit_fn=exchange_prefect_for_ray_future,
+            return_data=True,
+        )
+
+        return kwargs_ray_futures, upstream_ray_obj_refs
+
+    @staticmethod
+    def _run_prefect_task(func, *upstream_ray_obj_refs, **kwargs):
+        """Resolves Ray futures before calling the actual Prefect task function.
+
+        Passing upstream_ray_obj_refs directly as args enables Ray to wait for
+        upstream tasks before running this remote function.
+        This variable is otherwise unused as the ray object refs are also
+        contained in kwargs.
+        """
+
+        def resolve_ray_future(expr):
+            """Resolves Ray future."""
+            if isinstance(expr, ray.ObjectRef):
+                return ray.get(expr)
+            return expr
+
+        kwargs = visit_collection(kwargs, visit_fn=resolve_ray_future, return_data=True)
+
+        return func(**kwargs)
+
+    async def wait(self, key: UUID, timeout: float = None) -> Optional[State]:
+        ref = self._get_ray_ref(key)
+
+        result = None
+
+        with anyio.move_on_after(timeout):
+            # We await the reference directly instead of using `ray.get` so we can
+            # avoid blocking the event loop
+            try:
+                result = await ref
+            except RayTaskError as exc:
+                # unwrap the original exception that caused task failure, except for
+                # KeyboardInterrupt, which unwraps as TaskCancelledError
+                result = await exception_to_crashed_state(exc.cause)
+            except BaseException as exc:
+                result = await exception_to_crashed_state(exc)
+
+        return result
+
+    async def _start(self, exit_stack: AsyncExitStack):
+        """
+        Start the task runner and prep for context exit.
+
+        - Creates a cluster if an external address is not set.
+        - Creates a client to connect to the cluster.
+        - Pushes a call to wait for all running futures to complete on exit.
+        """
+        if self.address and self.address != "auto":
+            self.logger.info(
+                f"Connecting to an existing Ray instance at {self.address}"
+            )
+            init_args = (self.address,)
+        elif ray.is_initialized():
+            self.logger.info(
+                "Local Ray instance is already initialized. "
+                "Using existing local instance."
+            )
+            return
+        else:
+            self.logger.info("Creating a local Ray instance")
+            init_args = ()
+
+        context = ray.init(*init_args, **self.init_kwargs)
+        dashboard_url = getattr(context, "dashboard_url", None)
+        exit_stack.push(context)
+
+        # Display some information about the cluster
+        nodes = ray.nodes()
+        living_nodes = [node for node in nodes if node.get("alive")]
+        self.logger.info(f"Using Ray cluster with {len(living_nodes)} nodes.")
+
+        if dashboard_url:
+            self.logger.info(
+                f"The Ray UI is available at {dashboard_url}",
+            )
+
+    async def _shutdown_ray(self):
+        """
+        Shuts down the cluster.
+        """
+        self.logger.debug("Shutting down Ray cluster...")
+        ray.shutdown()
+
+    def _get_ray_ref(self, key: UUID) -> "ray.ObjectRef":
+        """
+        Retrieve the ray object reference corresponding to a prefect future.
+        """
+        return self._ray_refs[key]


Test patch:
diff --git a/scripts/generate_integration_package_tests_matrix.py b/scripts/generate_integration_package_tests_matrix.py
--- a/scripts/generate_integration_package_tests_matrix.py
+++ b/scripts/generate_integration_package_tests_matrix.py
@@ -11,6 +11,10 @@
     "3.12",
 ]
 
+SKIP_VERSIONS = {
+    "prefect-ray": ["3.8", "3.12"],
+}
+
 
 def get_changed_packages(commit_range: str) -> List[str]:
     # Get the list of changed files in the specified commit range
@@ -34,6 +38,8 @@ def generate_matrix(packages: List[str], python_versions: List[str]) -> Dict:
     matrix = {"include": []}
     for package in packages:
         for version in python_versions:
+            if version in SKIP_VERSIONS.get(package, []):
+                continue
             matrix["include"].append({"package": package, "python-version": version})
     return matrix
 
diff --git a/src/integrations/prefect-ray/tests/__init__.py b/src/integrations/prefect-ray/tests/__init__.py
new file mode 100644
diff --git a/src/integrations/prefect-ray/tests/conftest.py b/src/integrations/prefect-ray/tests/conftest.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-ray/tests/conftest.py
@@ -0,0 +1,16 @@
+import pytest
+
+from prefect.settings import PREFECT_ASYNC_FETCH_STATE_RESULT, temporary_settings
+from prefect.testing.utilities import prefect_test_harness
+
+
+@pytest.fixture(scope="session", autouse=True)
+def prefect_db():
+    with prefect_test_harness():
+        yield
+
+
+@pytest.fixture(autouse=True)
+def fetch_state_result():
+    with temporary_settings(updates={PREFECT_ASYNC_FETCH_STATE_RESULT: True}):
+        yield
diff --git a/src/integrations/prefect-ray/tests/test_block_standards.py b/src/integrations/prefect-ray/tests/test_block_standards.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-ray/tests/test_block_standards.py
@@ -0,0 +1,23 @@
+import pytest
+
+from prefect.blocks.core import Block
+from prefect.testing.standard_test_suites import BlockStandardTestSuite
+from prefect.utilities.dispatch import get_registry_for_type
+from prefect.utilities.importtools import to_qualified_name
+
+
+def find_module_blocks():
+    blocks = get_registry_for_type(Block)
+    module_blocks = [
+        block
+        for block in blocks.values()
+        if to_qualified_name(block).startswith("prefect_ray")
+    ]
+    return module_blocks
+
+
+@pytest.mark.parametrize("block", find_module_blocks())
+class TestAllBlocksAdhereToStandards(BlockStandardTestSuite):
+    @pytest.fixture
+    def block(self, block):
+        return block
diff --git a/src/integrations/prefect-ray/tests/test_context.py b/src/integrations/prefect-ray/tests/test_context.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-ray/tests/test_context.py
@@ -0,0 +1,33 @@
+from prefect_ray.context import RemoteOptionsContext, remote_options
+
+
+def test_remote_options_context():
+    input_remote_options = {"num_cpus": 4}
+    current_remote_options = RemoteOptionsContext(
+        current_remote_options=input_remote_options
+    ).current_remote_options
+    assert current_remote_options == input_remote_options
+
+
+def test_remote_options_context_get():
+    current_remote_options = RemoteOptionsContext.get().current_remote_options
+    assert current_remote_options == {}
+
+
+def test_remote_options():
+    with remote_options(num_cpus=4, num_gpus=None):
+        current_remote_options = RemoteOptionsContext.get().current_remote_options
+        assert current_remote_options == {"num_cpus": 4, "num_gpus": None}
+
+
+def test_remote_options_empty():
+    with remote_options():
+        current_remote_options = RemoteOptionsContext.get().current_remote_options
+        assert current_remote_options == {}
+
+
+def test_remote_options_override():
+    with remote_options(num_cpus=2, num_gpus=1):
+        with remote_options(num_cpus=4, num_gpus=2):
+            current_remote_options = RemoteOptionsContext.get().current_remote_options
+            assert current_remote_options == {"num_cpus": 4, "num_gpus": 2}
diff --git a/src/integrations/prefect-ray/tests/test_task_runners.py b/src/integrations/prefect-ray/tests/test_task_runners.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-ray/tests/test_task_runners.py
@@ -0,0 +1,357 @@
+import asyncio
+import logging
+import subprocess
+import sys
+import time
+import warnings
+from functools import partial
+from uuid import uuid4
+
+import pytest
+import ray
+import ray.cluster_utils
+from prefect_ray import RayTaskRunner
+from prefect_ray.context import remote_options
+from ray.exceptions import TaskCancelledError
+
+import prefect
+import prefect.engine
+import tests
+from prefect import flow, get_run_logger, task
+from prefect.states import State, StateType
+from prefect.tasks import TaskRun
+from prefect.testing.fixtures import (  # noqa: F401
+    hosted_api_server,
+    use_hosted_api_server,
+)
+from prefect.testing.standard_test_suites import TaskRunnerStandardTestSuite
+from prefect.testing.utilities import exceptions_equal
+
+
+@pytest.fixture(scope="session")
+def event_loop(request):
+    """
+    Redefine the event loop to support session/module-scoped fixtures;
+    see https://github.com/pytest-dev/pytest-asyncio/issues/68
+    When running on Windows we need to use a non-default loop for subprocess support.
+    """
+    if sys.platform == "win32" and sys.version_info >= (3, 8):
+        asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
+
+    policy = asyncio.get_event_loop_policy()
+
+    loop = policy.new_event_loop()
+
+    # configure asyncio logging to capture long running tasks
+    asyncio_logger = logging.getLogger("asyncio")
+    asyncio_logger.setLevel("WARNING")
+    asyncio_logger.addHandler(logging.StreamHandler())
+    loop.set_debug(True)
+    loop.slow_callback_duration = 0.25
+
+    try:
+        yield loop
+    finally:
+        loop.close()
+
+
+@pytest.fixture(scope="module")
+def machine_ray_instance():
+    """
+    Starts a ray instance for the current machine
+    """
+    subprocess.check_call(
+        [
+            "ray",
+            "start",
+            "--head",
+            "--include-dashboard",
+            "False",
+            "--disable-usage-stats",
+        ],
+        cwd=str(prefect.__development_base_path__),
+    )
+    try:
+        yield "ray://127.0.0.1:10001"
+    finally:
+        subprocess.run(["ray", "stop"])
+
+
+@pytest.fixture
+def default_ray_task_runner():
+    with warnings.catch_warnings():
+        # Ray does not properly close resources and we do not want their warnings to
+        # bubble into our test suite
+        # https://github.com/ray-project/ray/pull/22419
+        warnings.simplefilter("ignore", ResourceWarning)
+
+        yield RayTaskRunner()
+
+
+@pytest.fixture
+def ray_task_runner_with_existing_cluster(
+    machine_ray_instance,
+    use_hosted_api_server,  # noqa: F811
+    hosted_api_server,  # noqa: F811
+):
+    """
+    Generate a ray task runner that's connected to a ray instance running in a separate
+    process.
+
+    This tests connection via `ray://` which is a client-based connection.
+    """
+    yield RayTaskRunner(
+        address=machine_ray_instance,
+        init_kwargs={
+            "runtime_env": {
+                # Ship the 'tests' module to the workers or they will not be able to
+                # deserialize test tasks / flows
+                "py_modules": [tests]
+            }
+        },
+    )
+
+
+@pytest.fixture(scope="module")
+def inprocess_ray_cluster():
+    """
+    Starts a ray cluster in-process
+    """
+    cluster = ray.cluster_utils.Cluster(initialize_head=True)
+    try:
+        cluster.add_node()  # We need to add a second node for parallelism
+        yield cluster
+    finally:
+        cluster.shutdown()
+
+
+@pytest.fixture
+def ray_task_runner_with_inprocess_cluster(
+    inprocess_ray_cluster,
+    use_hosted_api_server,  # noqa: F811
+    hosted_api_server,  # noqa: F811
+):
+    """
+    Generate a ray task runner that's connected to an in-process cluster.
+
+    This tests connection via 'localhost' which is not a client-based connection.
+    """
+
+    yield RayTaskRunner(
+        address=inprocess_ray_cluster.address,
+        init_kwargs={
+            "runtime_env": {
+                # Ship the 'tests' module to the workers or they will not be able to
+                # deserialize test tasks / flows
+                "py_modules": [tests]
+            }
+        },
+    )
+
+
+@pytest.fixture
+def ray_task_runner_with_temporary_cluster(
+    use_hosted_api_server,  # noqa: F811
+    hosted_api_server,  # noqa: F811
+):
+    """
+    Generate a ray task runner that creates a temporary cluster.
+
+    This tests connection via 'localhost' which is not a client-based connection.
+    """
+
+    yield RayTaskRunner(
+        init_kwargs={
+            "runtime_env": {
+                # Ship the 'tests' module to the workers or they will not be able to
+                # deserialize test tasks / flows
+                "py_modules": [tests]
+            }
+        },
+    )
+
+
+task_runner_setups = [
+    default_ray_task_runner,
+    ray_task_runner_with_inprocess_cluster,
+    ray_task_runner_with_temporary_cluster,
+]
+
+if sys.version_info >= (3, 10):
+    task_runner_setups.append(ray_task_runner_with_existing_cluster)
+
+
+class TestRayTaskRunner(TaskRunnerStandardTestSuite):
+    @pytest.fixture(params=task_runner_setups)
+    def task_runner(self, request):
+        yield request.getfixturevalue(
+            request.param._pytestfixturefunction.name or request.param.__name__
+        )
+
+    def get_sleep_time(self) -> float:
+        """
+        Return an amount of time to sleep for concurrency tests.
+        The RayTaskRunner is prone to flaking on concurrency tests.
+        """
+        return 5.0
+
+    @pytest.mark.parametrize("exception", [KeyboardInterrupt(), ValueError("test")])
+    async def test_wait_captures_exceptions_as_crashed_state(
+        self, task_runner, exception
+    ):
+        """
+        Ray wraps the exception, interrupts will result in "Cancelled" tasks
+        or "Killed" workers while normal errors will result in a "RayTaskError".
+        We care more about the crash detection and
+        lack of re-raise here than the equality of the exception.
+        """
+
+        async def fake_orchestrate_task_run(task_run):
+            raise exception
+
+        task_run = TaskRun(
+            flow_run_id=uuid4(), task_key=str(uuid4()), dynamic_key="bar"
+        )
+
+        async with task_runner.start():
+            await task_runner.submit(
+                call=partial(fake_orchestrate_task_run, task_run=task_run),
+                key=task_run.id,
+            )
+
+            state = await task_runner.wait(task_run.id, 5)
+            assert state is not None, "wait timed out"
+            assert isinstance(state, State), "wait should return a state"
+            assert state.name == "Crashed"
+
+    @pytest.mark.parametrize(
+        "exceptions",
+        [
+            (KeyboardInterrupt(), TaskCancelledError),
+            (ValueError("test"), ValueError),
+        ],
+    )
+    async def test_exception_to_crashed_state_in_flow_run(
+        self, exceptions, task_runner, monkeypatch
+    ):
+        (raised_exception, state_exception_type) = exceptions
+
+        async def throws_exception_before_task_begins(
+            task, task_run, parameters, wait_for, result_factory, settings, **kwds
+        ):
+            """
+            Simulates an exception occurring while a remote task runner is attempting
+            to unpickle and run a Prefect task.
+            """
+            raise raised_exception
+
+        monkeypatch.setattr(
+            prefect.engine, "begin_task_run", throws_exception_before_task_begins
+        )
+
+        @task()
+        def test_task():
+            logger = get_run_logger()
+            logger.info("Ray should raise an exception before this task runs.")
+
+        @flow(task_runner=task_runner)
+        def test_flow():
+            future = test_task.submit()
+            future.wait(10)
+
+        # ensure that the type of exception raised by the flow matches the type of
+        # exception we expected the task runner to receive.
+        with pytest.raises(state_exception_type) as exc:
+            test_flow()
+            # If Ray passes the same exception type back, it should pass
+            # the equality check
+            if type(raised_exception) == state_exception_type:
+                assert exceptions_equal(raised_exception, exc)
+
+    def test_flow_and_subflow_both_with_task_runner(self, task_runner, tmp_file):
+        @task
+        def some_task(text):
+            tmp_file.write_text(text)
+
+        @flow(task_runner=RayTaskRunner())
+        def subflow():
+            some_task.submit("a")
+            some_task.submit("b")
+            some_task.submit("c")
+
+        @flow(task_runner=task_runner)
+        def base_flow():
+            subflow()
+            time.sleep(self.get_sleep_time())
+            some_task.submit("d")
+
+        base_flow()
+        assert tmp_file.read_text() == "d"
+
+    def test_ray_options(self):
+        @task
+        def process(x):
+            return x + 1
+
+        @flow(task_runner=RayTaskRunner())
+        def my_flow():
+            # equivalent to setting @ray.remote(max_calls=1)
+            with remote_options(max_calls=1):
+                process.submit(42)
+
+        my_flow()
+
+    def test_dependencies(self):
+        @task
+        def a():
+            time.sleep(self.get_sleep_time())
+
+        b = c = d = e = a
+
+        @flow(task_runner=RayTaskRunner())
+        def flow_with_dependent_tasks():
+            for _ in range(3):
+                a_future = a.submit(wait_for=[])
+                b_future = b.submit(wait_for=[a_future])
+
+                c.submit(wait_for=[b_future])
+                d.submit(wait_for=[b_future])
+                e.submit(wait_for=[b_future])
+
+        flow_with_dependent_tasks()
+
+    def test_sync_task_timeout(self, task_runner):
+        """
+        This test is inherited from the prefect testing module and it may not
+        appropriately skip on Windows. Here we skip it explicitly.
+        """
+        if sys.platform.startswith("win"):
+            pytest.skip("cancellation due to timeouts is not supported on Windows")
+        super().test_async_task_timeout(task_runner)
+
+    async def test_submit_and_wait(self, task_runner):
+        """
+        This test is inherited from the prefect testing module. The key difference
+        here is that task_runner is waiting longer than 5 seconds.
+        """
+        MAX_WAIT_TIME = 60
+
+        task_run = TaskRun(flow_run_id=uuid4(), task_key="foo", dynamic_key="bar")
+
+        async def fake_orchestrate_task_run(example_kwarg, task_run):
+            return State(
+                type=StateType.COMPLETED,
+                data=example_kwarg,
+            )
+
+        async with task_runner.start():
+            await task_runner.submit(
+                key=task_run.id,
+                call=partial(
+                    fake_orchestrate_task_run, task_run=task_run, example_kwarg=1
+                ),
+            )
+            state = await task_runner.wait(task_run.id, MAX_WAIT_TIME)
+            assert state is not None, "wait timed out"
+            assert isinstance(state, State), "wait should return a state"
+            assert await state.result() == 1
diff --git a/src/integrations/prefect-ray/tests/test_version.py b/src/integrations/prefect-ray/tests/test_version.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-ray/tests/test_version.py
@@ -0,0 +1,9 @@
+from packaging.version import Version
+
+
+def test_version():
+    from prefect_ray import __version__
+
+    assert isinstance(__version__, str)
+    assert Version(__version__)
+    assert __version__.startswith("0.")