perf(tests): reduce runtime and GPU memory

.github/workflows/main.yml

@@ -150,20 +150,35 @@ jobs:
       - name: Verify Scene Engine gensim installation
         run: |
           python -c "import matplotlib, numpy, open3d, requests, scipy, shapely, trimesh; from PIL import Image; import embodichain.gen_sim.scene_engine.pipeline.generate"
-          pytest tests/gen_sim/scene_engine -q
 
       - name: Run default tests
         run: |
-          echo "Default test suite (GPU-marked tests are skipped)"
+          echo "Documentation tests"
           export HF_ENDPOINT=https://hf-mirror.com
           pytest tests/docs -q --confcutdir=tests/docs
-          pytest tests --ignore=tests/docs
 
-      - name: Run GPU tests serially
+          echo "Non-simulation tests (four workers, CUDA available for native imports)"
+          pytest tests --ignore=tests/docs \
+            -m "not requires_sim and not gpu" -n 4 --dist loadgroup
+
+      - name: Run real-simulation tests serially
+        run: |
+          echo "Real-simulation tests (one process, GPU tests excluded)"
+          export HF_ENDPOINT=https://hf-mirror.com
+          pytest tests --ignore=tests/docs -m "requires_sim and not gpu"
+
+      - name: Run distributed GPU tests in isolation
+        run: |
+          echo "Distributed GPU tests (isolated CUDA process tree)"
+          export HF_ENDPOINT=https://hf-mirror.com
+          pytest tests/learning/test_rl_distributed.py --run-gpu -m gpu
+
+      - name: Run remaining GPU tests serially
         run: |
           echo "Dedicated GPU test suite"
           export HF_ENDPOINT=https://hf-mirror.com
-          pytest tests --run-gpu -m gpu
+          pytest tests --ignore=tests/docs \
+            --ignore=tests/learning/test_rl_distributed.py --run-gpu -m gpu
 
   release-build:
     if: startsWith(github.ref, 'refs/tags/v')

embodichain/lab/sim/objects/articulation.py

@@ -1683,23 +1683,27 @@ def set_joint_drive(
         cache_env_ids = self._resolve_env_ids(env_ids)
         cache_joint_ids = self._resolve_joint_ids(joint_ids)
 
+        def _drive_arg(value: torch.Tensor, index: int) -> float | np.ndarray:
+            result = value[index].detach().cpu().numpy()
+            return result.item() if result.size == 1 else result
+
         for i, env_idx in enumerate(local_env_ids):
             drive_args = {
                 "drive_type": get_dexsim_drive_type(drive_type),
                 "joint_ids": local_joint_ids,
             }
             if stiffness is not None:
-                drive_args["stiffness"] = stiffness[i].cpu().numpy()
+                drive_args["stiffness"] = _drive_arg(stiffness, i)
             if damping is not None:
-                drive_args["damping"] = damping[i].cpu().numpy()
+                drive_args["damping"] = _drive_arg(damping, i)
             if max_effort is not None:
-                drive_args["max_force"] = max_effort[i].cpu().numpy()
+                drive_args["max_force"] = _drive_arg(max_effort, i)
             if max_velocity is not None:
-                drive_args["max_velocity"] = max_velocity[i].cpu().numpy()
+                drive_args["max_velocity"] = _drive_arg(max_velocity, i)
             if friction is not None:
-                drive_args["joint_friction"] = friction[i].cpu().numpy()
+                drive_args["joint_friction"] = _drive_arg(friction, i)
             if armature is not None:
-                drive_args["armature"] = armature[i].cpu().numpy()
+                drive_args["armature"] = _drive_arg(armature, i)
             self._entities[env_idx].set_drive(**drive_args)
 
         if max_velocity is not None:

embodichain/lab/sim/sim_manager.py

@@ -3151,19 +3151,33 @@ def _sever_wrapper_refs(obj_registry):
         gc.collect()
 
     @staticmethod
-    def flush_cleanup_queue():
-        """Dequeue executor and synchronization barrier provided for top-level main loop / Pytest Fixture calls"""
+    def flush_cleanup_queue() -> None:
+        """Run pending destruction tasks and wait for their scenes to disappear.
+
+        An empty queue means that no manager requested destruction.  In that
+        case, returning immediately is important: other managers may still own
+        live worlds, and waiting for the global world count to reach zero would
+        block until the timeout even though there is nothing to clean up.
+        """
         import gc
 
-        while not SimulationManager._cleanup_queue.empty():
-            task = SimulationManager._cleanup_queue.get_nowait()
+        drained_task = False
+        while True:
+            try:
+                task = SimulationManager._cleanup_queue.get_nowait()
+            except queue.Empty:
+                break
+
+            drained_task = True
             try:
                 task()
             except Exception as e:
                 from embodichain.utils import logger
 
                 logger.log_error(f"Error during delayed destruction: {e}")
-                pass
+
+        if not drained_task:
+            return
 
         # After the queue is emptied, perform a top-level full GC to thoroughly reclaim dead objects that haven't released their RefPtrs yet
         gc.collect()

embodichain/lab/sim/solvers/srs_solver.py

@@ -14,6 +14,8 @@
 # limitations under the License.
 # ----------------------------------------------------------------------------
 
+from __future__ import annotations
+
 import torch
 import numpy as np
 import warp as wp
@@ -36,7 +38,7 @@
     from embodichain.lab.sim.robots.dexforce_w1.params import W1ArmKineParams
 
 
-all = ["SRSSolver", "SRSSolverCfg"]
+__all__ = ["SRSSolver", "SRSSolverCfg"]
 
 
 @configclass
@@ -404,13 +406,16 @@ def _process_single_solution(
             return success_tensor, ik_qpos_tensor[:, :1, :]
 
     def _get_each_ik(
-        self, target_pose: np.ndarray, nsparam: float, config: np.ndarray
+        self,
+        target_pose: np.ndarray | torch.Tensor,
+        nsparam: float,
+        config: np.ndarray,
     ) -> tuple[bool, np.ndarray | None]:
         """
         Computes the inverse kinematics for a given target pose, normalization parameter, and configuration.
 
         Args:
-            target_pose (np.ndarray): 4x4 target pose matrix.
+            target_pose (np.ndarray | torch.Tensor): 4x4 target pose matrix.
             nsparam (float): Normalization parameter (angle).
             config (np.ndarray): Configuration index.
 
@@ -419,7 +424,10 @@ def _get_each_ik(
             np.ndarray: List of joint solutions (7) or None if no solution is found.
         """
         # Validate the target pose matrix
-        target_pose = np.array(target_pose)
+        if isinstance(target_pose, torch.Tensor):
+            target_pose = target_pose.detach().cpu().numpy()
+        else:
+            target_pose = np.asarray(target_pose)
         if target_pose.ndim == 3 and target_pose.shape[0] == 1:
             target_pose = target_pose[0]  # Extract the first matrix
         if target_pose.shape != (4, 4):

pyproject.toml

@@ -106,9 +106,16 @@ include-package-data = false
 [tool.black]
 
 [tool.pytest.ini_options]
+addopts = ["-m", "not slow"]
+filterwarnings = [
+    "ignore:`torch\\.jit\\.script(_method)?` is deprecated.*:DeprecationWarning:torch\\.jit\\._script",
+]
 markers = [
     "slow: marks tests as slow (deselect with '-m \"not slow\"')",
     "requires_sim: marks tests that require a real simulation backend",
+    "no_sim: marks mock-only tests that must not initialize a simulation backend",
+    "requires_tasks: marks tests that require installed task-package discovery",
     "gpu: marks tests that execute CUDA/GPU code (run serially)",
     "renderer: marks tests that exercise a rendering backend",
+    "xdist_group(name): keeps resource-sharing tests on one pytest-xdist worker",
 ]

tests/conftest.py

@@ -16,14 +16,17 @@
 
 from __future__ import annotations
 
+from functools import lru_cache
+
 import inspect
 import os
+import re
 import pytest
 
 os.environ.setdefault("EMBODICHAIN_SIM_EXIT_PROCESS", "0")
 
 
-@pytest.fixture(scope="session", autouse=True)
+@pytest.fixture(scope="session")
 def _discover_task_packages():
     """Discover all installed task packages once per test session.
 
@@ -45,6 +48,14 @@ def _discover_task_packages():
     execute_init_hooks()
 
 
+@pytest.fixture(autouse=True)
+def _discover_task_packages_for_marked_tests(request):
+    """Discover task packages only for tests that consume registered tasks."""
+    if request.node.get_closest_marker("requires_tasks") is not None:
+        request.getfixturevalue("_discover_task_packages")
+    yield
+
+
 def pytest_addoption(parser):
     parser.addoption(
         "--renderer",
@@ -68,34 +79,78 @@ def pytest_configure(config):
                 f"Invalid renderer: {renderer}. Must be one of 'hybrid', 'fast-rt'"
             )
 
-        # Override the global default renderer in the simulation config
-        from embodichain.lab.sim import cfg
-
-        cfg.DEFAULT_RENDERER = renderer
-
         # DexSim initialization is intentionally deferred to the first real-simulation
         # test.  Most of the suite consists of pure-Python tests and should not acquire
         # a CUDA/Vulkan context merely because pytest has started.
 
 
+_SIMULATION_MANAGER_CONSTRUCTOR = re.compile(r"\bSimulationManager\s*\(")
+
+
+@lru_cache(maxsize=None)
+def _source_constructs_real_sim(obj):
+    """Return whether an object's source directly constructs a simulation manager."""
+    try:
+        return (
+            _SIMULATION_MANAGER_CONSTRUCTOR.search(inspect.getsource(obj)) is not None
+        )
+    except (OSError, TypeError):
+        return False
+
+
+def _callable_constructs_real_sim(obj):
+    """Check a callable and the module-level helpers it directly references."""
+    if _source_constructs_real_sim(obj):
+        return True
+
+    code = getattr(obj, "__code__", None)
+    module = inspect.getmodule(obj)
+    if code is None or module is None:
+        return False
+
+    for name in code.co_names:
+        helper = vars(module).get(name)
+        if inspect.isfunction(helper) and _source_constructs_real_sim(helper):
+            return True
+    return False
+
+
 def _requires_real_sim(item):
-    """Return whether a test module creates a real SimulationManager."""
+    """Return whether a test item creates a real SimulationManager."""
     if item.get_closest_marker("requires_sim") is not None:
         return True
+    if item.get_closest_marker("no_sim") is not None:
+        return False
+
     module = getattr(item, "module", None)
     if module is not None and "SimulationManager" in vars(module):
         return True
 
-    # Some planner regression tests intentionally import the simulation manager
-    # inside the test body to keep their module import lightweight.
-    try:
-        return "SimulationManager" in inspect.getsource(item.obj)
-    except (OSError, TypeError):
-        return False
+    if _callable_constructs_real_sim(item.obj):
+        return True
+
+    test_class = getattr(item, "cls", None)
+    if test_class is not None:
+        for cls in test_class.__mro__:
+            if _source_constructs_real_sim(cls):
+                return True
+
+    fixture_info = getattr(item, "_fixtureinfo", None)
+    fixture_defs_by_name = getattr(fixture_info, "name2fixturedefs", {})
+    for fixture_defs in fixture_defs_by_name.values():
+        for fixture_def in fixture_defs:
+            if _callable_constructs_real_sim(fixture_def.func):
+                return True
+
+    return False
 
 
 def _initialize_sim_engine(renderer):
     """Initialize DexSim once, immediately before the first real-sim test."""
+    from embodichain.lab.sim import cfg
+
+    cfg.DEFAULT_RENDERER = renderer
+
     import dexsim
     import dexsim.types
 
@@ -113,13 +168,20 @@ def _initialize_sim_engine(renderer):
     dexsim.init_sim_engine(sim_config)
 
 
+@pytest.hookimpl(tryfirst=True)
 def pytest_collection_modifyitems(config, items):
     """Classify real-simulation tests for fast and resource-aware test selection."""
     for item in items:
         nodeid = item.nodeid.lower()
         requires_sim = _requires_real_sim(item)
         if requires_sim:
             item.add_marker(pytest.mark.requires_sim)
+            item.add_marker(pytest.mark.xdist_group(name="simulation"))
+        else:
+            module_name = getattr(item.module, "__name__", "unknown")
+            item.add_marker(
+                pytest.mark.xdist_group(name=f"module-{module_name.replace('.', '-')}")
+            )
         if "cuda" in nodeid or "gpu" in nodeid:
             item.add_marker(pytest.mark.gpu)
         if requires_sim and (

tests/data_pipeline/test_online_data.py

@@ -120,7 +120,10 @@ def _make_fake_engine(
     engine._close_signal = engine._mp_ctx.Event()
     engine._sample_count = engine._mp_ctx.Value("i", 0)
 
-    engine.start()
+    # Sampling tests intentionally bypass the simulation worker.  Starting it
+    # here would exercise unrelated environment setup and make teardown wait
+    # for the subprocess timeout when the synthetic config cannot run a task.
+    engine._sim_process = None
 
     return engine
 

tests/gym/envs/test_base_env.py

@@ -37,6 +37,8 @@
 
 NUM_ENVS = 10
 
+pytestmark = pytest.mark.requires_sim
+
 
 @register_env("RandomReach-v1", max_episode_steps=100, override=True)
 class RandomReachEnv(BaseEnv):

tests/gym/envs/test_embodied_env.py

@@ -39,6 +39,8 @@
 
 NUM_ENVS = 2
 
+pytestmark = pytest.mark.requires_sim
+
 urdf_path = get_data_path("UniversalRobots/UR5/UR5.urdf")
 METADATA = {
     "id": "EmbodiedEnv-v1",

tests/gym/envs/test_profiler_integration.py

@@ -34,6 +34,8 @@
 from embodichain.lab.sim.cfg import RobotCfg
 from embodichain.lab.sim.objects import Robot
 
+pytestmark = pytest.mark.requires_sim
+
 
 @register_env("ProfilerProbe-v1", max_episode_steps=100, override=True)
 class _ProfilerProbeEnv(BaseEnv):

tests/gym/envs/test_replay.py

@@ -21,6 +21,7 @@
 from unittest.mock import MagicMock, patch
 
 import numpy as np
+import pytest
 import torch
 
 from embodichain.data import get_data_path
@@ -33,6 +34,8 @@
 from embodichain.lab.gym.envs.managers.actions import DeltaQposTerm
 from embodichain.lab.gym.envs.managers.cfg import ActionTermCfg
 
+pytestmark = [pytest.mark.requires_sim, pytest.mark.slow]
+
 
 @register_env("ReplayTest-v1", max_episode_steps=100, override=True)
 class ReplayTestEnv(EmbodiedEnv):

tests/learning/test_newton_planar_reach.py

@@ -351,6 +351,7 @@ def test_evaluation_uses_eval_mode_and_restores_policy_mode(
     assert policy.training is initial_training_mode
 
 
+@pytest.mark.slow
 def test_apg_training_generalizes_to_held_out_reaches() -> None:
     result = train_planar_reach(
         NewtonPlanarReachTrainingCfg(