Wait for upstream object refs before starting task

prefect_ray/task_runners.py

@@ -145,7 +145,9 @@ async def submit(
                 "The task runner must be started before submitting work."
             )
 
-        call_kwargs = self._exchange_prefect_for_ray_futures(call.keywords)
+        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
@@ -154,18 +156,22 @@ async def submit(
             ray_decorator = ray.remote(**remote_options)
         else:
             ray_decorator = ray.remote
+
         self._ray_refs[key] = ray_decorator(self._run_prefect_task).remote(
-            sync_compatible(call.func), **call_kwargs
+            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
 
@@ -175,11 +181,17 @@ def exchange_prefect_for_ray_future(expr):
             return_data=True,
         )
 
-        return kwargs_ray_futures
+        return kwargs_ray_futures, upstream_ray_obj_refs
 
     @staticmethod
-    def _run_prefect_task(func, *args, **kwargs):
-        """Resolves Ray futures before calling the actual Prefect task function."""
+    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."""
@@ -189,7 +201,7 @@ def resolve_ray_future(expr):
 
         kwargs = visit_collection(kwargs, visit_fn=resolve_ray_future, return_data=True)
 
-        return func(*args, **kwargs)
+        return func(**kwargs)
 
     async def wait(self, key: UUID, timeout: float = None) -> Optional[State]:
         ref = self._get_ray_ref(key)