Use public api from ValueMesh

src/forge/actors/policy.py

@@ -248,7 +248,7 @@ async def setup(self):
         # Setup scheduler
         # TODO: Add support for `log_stats`
         kv_cache_configs = await self.policy_worker.setup_kv_cache.call()
-        kv_cache_config = kv_cache_configs._values[0]
+        _, kv_cache_config = next(kv_cache_configs.items())
         self.vllm_config.cache_config.num_gpu_blocks = kv_cache_config.num_blocks
         self.vllm_config.cache_config.num_cpu_blocks = 0
 

src/forge/controller/service/replica.py

@@ -242,12 +242,8 @@ async def _process_single_request(self, request: ServiceRequest) -> bool:
             try:
                 result = await endpoint_func.call(*request.args, **request.kwargs)
                 # Unwrap ValueMesh if configured to return first rank result
-                if (
-                    self.return_first_rank_result
-                    and hasattr(result, "_values")
-                    and result._values
-                ):
-                    result = result._values[0]
+                if self.return_first_rank_result:
+                    _, result = next(result.items())
                 request.future.set_result(result)
             except ActorError as e:
                 logger.warning(f"Got failure on replica {self.idx}. Error:\n{e}")