Update

[ghstack-poisoned]

Author: vmoens

torchrl/modules/inference_server/_monarch.py

@@ -29,13 +29,22 @@ def __init__(self, client: _MonarchInferenceClient, req_id: int):
         self._req_id = req_id
         self._result: Any = _SENTINEL
 
+    def done(self) -> bool:
+        """Return ``True`` if the result is available without blocking."""
+        if self._result is not _SENTINEL:
+            return True
+        try:
+            self._result = self._client._get_result(self._req_id, timeout=0)
+        except queue.Empty:
+            return False
+        return True
+
     def result(self, timeout: float | None = None) -> TensorDictBase:
         """Block until the result is available."""
         if self._result is _SENTINEL:
-            item = self._client._get_result(self._req_id, timeout=timeout)
-            if isinstance(item, BaseException):
-                raise item
-            self._result = item
+            self._result = self._client._get_result(self._req_id, timeout=timeout)
+        if isinstance(self._result, BaseException):
+            raise self._result
         return self._result
 
 

torchrl/modules/inference_server/_mp.py

@@ -33,6 +33,16 @@ def __init__(self, client: _MPInferenceClient, req_id: int):
         self._req_id = req_id
         self._result: Any = _SENTINEL
 
+    def done(self) -> bool:
+        """Return ``True`` if the result is available without blocking."""
+        if self._result is not _SENTINEL:
+            return True
+        try:
+            self._result = self._client._get_result(self._req_id, timeout=0)
+        except queue.Empty:
+            return False
+        return True
+
     def result(self, timeout: float | None = None) -> TensorDictBase:
         """Block until the result is available.
 
@@ -44,10 +54,9 @@ def result(self, timeout: float | None = None) -> TensorDictBase:
             Exception: if the server set an exception instead of a result.
         """
         if self._result is _SENTINEL:
-            item = self._client._get_result(self._req_id, timeout=timeout)
-            if isinstance(item, BaseException):
-                raise item
-            self._result = item
+            self._result = self._client._get_result(self._req_id, timeout=timeout)
+        if isinstance(self._result, BaseException):
+            raise self._result
         return self._result
 
 

torchrl/modules/inference_server/_ray.py

@@ -32,13 +32,22 @@ def __init__(self, client: _RayInferenceClient, req_id: int):
         self._req_id = req_id
         self._result: Any = _SENTINEL
 
+    def done(self) -> bool:
+        """Return ``True`` if the result is available without blocking."""
+        if self._result is not _SENTINEL:
+            return True
+        try:
+            self._result = self._client._get_result(self._req_id, timeout=0)
+        except queue.Empty:
+            return False
+        return True
+
     def result(self, timeout: float | None = None) -> TensorDictBase:
         """Block until the result is available."""
         if self._result is _SENTINEL:
-            item = self._client._get_result(self._req_id, timeout=timeout)
-            if isinstance(item, BaseException):
-                raise item
-            self._result = item
+            self._result = self._client._get_result(self._req_id, timeout=timeout)
+        if isinstance(self._result, BaseException):
+            raise self._result
         return self._result