From 1cba47da20482f81016e162925063180da8dbbf6 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Wed, 24 Sep 2025 14:49:37 +0200 Subject: [PATCH] chore(async): move async related code to its directory at top level (#2003) * chore(async): move async related code to its directory at top level * chore(style): apply pre-commit to renamed headers * test(async): fix async imports * docs(async): update async headers doc --- docs/source/async.mdx | 16 ++++++++-------- .../server => async_inference}/configs.py | 3 ++- .../server => async_inference}/constants.py | 0 .../server => async_inference}/helpers.py | 0 .../policy_server.py | 19 ++++++++++--------- .../robot_client.py | 19 ++++++++++--------- tests/async_inference/test_e2e.py | 8 ++++---- tests/async_inference/test_helpers.py | 4 ++-- tests/async_inference/test_policy_server.py | 8 ++++---- tests/async_inference/test_robot_client.py | 8 ++++---- 10 files changed, 44 insertions(+), 41 deletions(-) rename src/lerobot/{scripts/server => async_inference}/configs.py (99%) rename src/lerobot/{scripts/server => async_inference}/constants.py (100%) rename src/lerobot/{scripts/server => async_inference}/helpers.py (100%) rename src/lerobot/{scripts/server => async_inference}/policy_server.py (98%) rename src/lerobot/{scripts/server => async_inference}/robot_client.py (98%) diff --git a/docs/source/async.mdx b/docs/source/async.mdx index 397c513c..c66cdb14 100644 --- a/docs/source/async.mdx +++ b/docs/source/async.mdx @@ -31,7 +31,7 @@ Then, spin up a policy server (in one terminal, or in a separate machine) specif You can spin up a policy server running: ```shell -python src/lerobot/scripts/server/policy_server.py \ +python src/lerobot/async_inference/policy_server.py \ --host=127.0.0.1 \ --port=8080 \ ``` @@ -39,7 +39,7 @@ python src/lerobot/scripts/server/policy_server.py \ This will start a policy server listening on `127.0.0.1:8080` (`localhost`, port 8080). At this stage, the policy server is empty, as all information related to which policy to run and with which parameters are specified during the first handshake with the client. Spin up a client with: ```shell -python src/lerobot/scripts/server/robot_client.py \ +python src/lerobot/async_inference/robot_client.py \ --server_address=127.0.0.1:8080 \ # SERVER: the host address and port of the policy server --robot.type=so100_follower \ # ROBOT: your robot type --robot.port=/dev/tty.usbmodem585A0076841 \ # ROBOT: your robot port @@ -122,8 +122,8 @@ python -m lerobot.scripts.server.policy_server \ ```python -from lerobot.scripts.server.configs import PolicyServerConfig -from lerobot.scripts.server.policy_server import serve +from lerobot.async_inference.configs import PolicyServerConfig +from lerobot.async_inference.policy_server import serve config = PolicyServerConfig( host="localhost", @@ -148,7 +148,7 @@ The `RobotClient` streams observations to the `PolicyServer`, and receives actio ```bash -python src/lerobot/scripts/server/robot_client.py \ +python src/lerobot/async_inference/robot_client.py \ --server_address=127.0.0.1:8080 \ # SERVER: the host address and port of the policy server --robot.type=so100_follower \ # ROBOT: your robot type --robot.port=/dev/tty.usbmodem585A0076841 \ # ROBOT: your robot port @@ -171,9 +171,9 @@ python src/lerobot/scripts/server/robot_client.py \ import threading from lerobot.robots.so100_follower import SO100FollowerConfig from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig -from lerobot.scripts.server.configs import RobotClientConfig -from lerobot.scripts.server.robot_client import RobotClient -from lerobot.scripts.server.helpers import visualize_action_queue_size +from lerobot.async_inference.configs import RobotClientConfig +from lerobot.async_inference.robot_client import RobotClient +from lerobot.async_inference.helpers import visualize_action_queue_size # 1. Create the robot instance """Check out the cameras available in your setup by running `python lerobot/find_cameras.py`""" diff --git a/src/lerobot/scripts/server/configs.py b/src/lerobot/async_inference/configs.py similarity index 99% rename from src/lerobot/scripts/server/configs.py rename to src/lerobot/async_inference/configs.py index 5be46485..24f889df 100644 --- a/src/lerobot/scripts/server/configs.py +++ b/src/lerobot/async_inference/configs.py @@ -18,7 +18,8 @@ from dataclasses import dataclass, field import torch from lerobot.robots.config import RobotConfig -from lerobot.scripts.server.constants import ( + +from .constants import ( DEFAULT_FPS, DEFAULT_INFERENCE_LATENCY, DEFAULT_OBS_QUEUE_TIMEOUT, diff --git a/src/lerobot/scripts/server/constants.py b/src/lerobot/async_inference/constants.py similarity index 100% rename from src/lerobot/scripts/server/constants.py rename to src/lerobot/async_inference/constants.py diff --git a/src/lerobot/scripts/server/helpers.py b/src/lerobot/async_inference/helpers.py similarity index 100% rename from src/lerobot/scripts/server/helpers.py rename to src/lerobot/async_inference/helpers.py diff --git a/src/lerobot/scripts/server/policy_server.py b/src/lerobot/async_inference/policy_server.py similarity index 98% rename from src/lerobot/scripts/server/policy_server.py rename to src/lerobot/async_inference/policy_server.py index 0ed446d3..12572706 100644 --- a/src/lerobot/scripts/server/policy_server.py +++ b/src/lerobot/async_inference/policy_server.py @@ -15,7 +15,7 @@ """ Example: ```shell -python src/lerobot/scripts/server/policy_server.py \ +python src/lerobot/async_inference/policy_server.py \ --host=127.0.0.1 \ --port=8080 \ --fps=30 \ @@ -38,9 +38,15 @@ import grpc import torch from lerobot.policies.factory import get_policy_class -from lerobot.scripts.server.configs import PolicyServerConfig -from lerobot.scripts.server.constants import SUPPORTED_POLICIES -from lerobot.scripts.server.helpers import ( +from lerobot.transport import ( + services_pb2, # type: ignore + services_pb2_grpc, # type: ignore +) +from lerobot.transport.utils import receive_bytes_in_chunks + +from .configs import PolicyServerConfig +from .constants import SUPPORTED_POLICIES +from .helpers import ( FPSTracker, Observation, RemotePolicyConfig, @@ -50,11 +56,6 @@ from lerobot.scripts.server.helpers import ( observations_similar, raw_observation_to_observation, ) -from lerobot.transport import ( - services_pb2, # type: ignore - services_pb2_grpc, # type: ignore -) -from lerobot.transport.utils import receive_bytes_in_chunks class PolicyServer(services_pb2_grpc.AsyncInferenceServicer): diff --git a/src/lerobot/scripts/server/robot_client.py b/src/lerobot/async_inference/robot_client.py similarity index 98% rename from src/lerobot/scripts/server/robot_client.py rename to src/lerobot/async_inference/robot_client.py index 939d5cea..c969bc60 100644 --- a/src/lerobot/scripts/server/robot_client.py +++ b/src/lerobot/async_inference/robot_client.py @@ -15,7 +15,7 @@ """ Example command: ```shell -python src/lerobot/scripts/server/robot_client.py \ +python src/lerobot/async_inference/robot_client.py \ --robot.type=so100_follower \ --robot.port=/dev/tty.usbmodem58760431541 \ --robot.cameras="{ front: {type: opencv, index_or_path: 0, width: 1920, height: 1080, fps: 30}}" \ @@ -57,9 +57,15 @@ from lerobot.robots import ( # noqa: F401 so100_follower, so101_follower, ) -from lerobot.scripts.server.configs import RobotClientConfig -from lerobot.scripts.server.constants import SUPPORTED_ROBOTS -from lerobot.scripts.server.helpers import ( +from lerobot.transport import ( + services_pb2, # type: ignore + services_pb2_grpc, # type: ignore +) +from lerobot.transport.utils import grpc_channel_options, send_bytes_in_chunks + +from .configs import RobotClientConfig +from .constants import SUPPORTED_ROBOTS +from .helpers import ( Action, FPSTracker, Observation, @@ -72,11 +78,6 @@ from lerobot.scripts.server.helpers import ( validate_robot_cameras_for_policy, visualize_action_queue_size, ) -from lerobot.transport import ( - services_pb2, # type: ignore - services_pb2_grpc, # type: ignore -) -from lerobot.transport.utils import grpc_channel_options, send_bytes_in_chunks class RobotClient: diff --git a/tests/async_inference/test_e2e.py b/tests/async_inference/test_e2e.py index 1c0400e6..2689f061 100644 --- a/tests/async_inference/test_e2e.py +++ b/tests/async_inference/test_e2e.py @@ -48,11 +48,11 @@ def test_async_inference_e2e(monkeypatch): # Import grpc-dependent modules inside the test function import grpc + from lerobot.async_inference.configs import PolicyServerConfig, RobotClientConfig + from lerobot.async_inference.helpers import map_robot_keys_to_lerobot_features + from lerobot.async_inference.policy_server import PolicyServer + from lerobot.async_inference.robot_client import RobotClient from lerobot.robots.utils import make_robot_from_config - from lerobot.scripts.server.configs import PolicyServerConfig, RobotClientConfig - from lerobot.scripts.server.helpers import map_robot_keys_to_lerobot_features - from lerobot.scripts.server.policy_server import PolicyServer - from lerobot.scripts.server.robot_client import RobotClient from lerobot.transport import ( services_pb2, # type: ignore services_pb2_grpc, # type: ignore diff --git a/tests/async_inference/test_helpers.py b/tests/async_inference/test_helpers.py index e0b79737..f1c7636e 100644 --- a/tests/async_inference/test_helpers.py +++ b/tests/async_inference/test_helpers.py @@ -19,8 +19,7 @@ import time import numpy as np import torch -from lerobot.configs.types import FeatureType, PolicyFeature -from lerobot.scripts.server.helpers import ( +from lerobot.async_inference.helpers import ( FPSTracker, TimedAction, TimedObservation, @@ -30,6 +29,7 @@ from lerobot.scripts.server.helpers import ( raw_observation_to_observation, resize_robot_observation_image, ) +from lerobot.configs.types import FeatureType, PolicyFeature # --------------------------------------------------------------------- # FPSTracker diff --git a/tests/async_inference/test_policy_server.py b/tests/async_inference/test_policy_server.py index 5c795e7e..c5c52460 100644 --- a/tests/async_inference/test_policy_server.py +++ b/tests/async_inference/test_policy_server.py @@ -65,8 +65,8 @@ class MockPolicy: def policy_server(): """Fresh `PolicyServer` instance with a stubbed-out policy model.""" # Import only when the test actually runs (after decorator check) - from lerobot.scripts.server.configs import PolicyServerConfig - from lerobot.scripts.server.policy_server import PolicyServer + from lerobot.async_inference.configs import PolicyServerConfig + from lerobot.async_inference.policy_server import PolicyServer test_config = PolicyServerConfig(host="localhost", port=9999) server = PolicyServer(test_config) @@ -95,7 +95,7 @@ def policy_server(): def _make_obs(state: torch.Tensor, timestep: int = 0, must_go: bool = False): """Create a TimedObservation with a given state vector.""" # Import only when needed - from lerobot.scripts.server.helpers import TimedObservation + from lerobot.async_inference.helpers import TimedObservation return TimedObservation( observation={ @@ -191,7 +191,7 @@ def test_obs_sanity_checks(policy_server): def test_predict_action_chunk(monkeypatch, policy_server): """End-to-end test of `_predict_action_chunk` with a stubbed _get_action_chunk.""" # Import only when needed - from lerobot.scripts.server.policy_server import PolicyServer + from lerobot.async_inference.policy_server import PolicyServer # Force server to act-style policy; patch method to return deterministic tensor policy_server.policy_type = "act" diff --git a/tests/async_inference/test_robot_client.py b/tests/async_inference/test_robot_client.py index 51db2c3a..dfdb8ce4 100644 --- a/tests/async_inference/test_robot_client.py +++ b/tests/async_inference/test_robot_client.py @@ -38,8 +38,8 @@ def robot_client(): """Fresh `RobotClient` instance for each test case (no threads started). Uses DummyRobot.""" # Import only when the test actually runs (after decorator check) - from lerobot.scripts.server.configs import RobotClientConfig - from lerobot.scripts.server.robot_client import RobotClient + from lerobot.async_inference.configs import RobotClientConfig + from lerobot.async_inference.robot_client import RobotClient from tests.mocks.mock_robot import MockRobotConfig test_config = MockRobotConfig() @@ -73,7 +73,7 @@ def robot_client(): def _make_actions(start_ts: float, start_t: int, count: int): """Generate `count` consecutive TimedAction objects starting at timestep `start_t`.""" - from lerobot.scripts.server.helpers import TimedAction + from lerobot.async_inference.helpers import TimedAction fps = 30 # emulates most common frame-rate actions = [] @@ -124,7 +124,7 @@ def test_aggregate_action_queues_combines_actions_in_overlap( ): """`_aggregate_action_queues` must combine actions on overlapping timesteps according to the provided aggregate_fn, here tested with multiple coefficients.""" - from lerobot.scripts.server.helpers import TimedAction + from lerobot.async_inference.helpers import TimedAction robot_client.chunks_received = 0