[pydrake] Fix diagram memory leaks (RobotLocomotion#22221)

bindings/pydrake/common/ref_cycle_pybind.h

@@ -10,8 +10,8 @@ namespace internal {
 
 /* pydrake::internal::ref_cycle is a custom call policy for pybind11.
 
-   For an overview of other call policies, See
-   https://pybind11.readthedocs.io/en/stable/advanced/functions.html#additional-call-policies
+  For an overview of other call policies, See
+  https://pybind11.readthedocs.io/en/stable/advanced/functions.html#additional-call-policies
 
   `ref_cycle` creates a reference count cycle that Python's cyclic garbage
   collection can find and collect, once the cycle's objects are no longer

bindings/pydrake/geometry/BUILD.bazel

@@ -38,11 +38,13 @@ drake_pybind_library(
         "//bindings/pydrake/common:default_scalars_pybind",
         "//bindings/pydrake/common:deprecation_pybind",
         "//bindings/pydrake/common:identifier_pybind",
+        "//bindings/pydrake/common:ref_cycle_pybind",
         "//bindings/pydrake/common:serialize_pybind",
         "//bindings/pydrake/common:sorted_pair_pybind",
         "//bindings/pydrake/common:type_pack",
         "//bindings/pydrake/common:type_safe_index_pybind",
         "//bindings/pydrake/common:value_pybind",
+        "//bindings/pydrake/systems:builder_life_support_pybind",
     ],
     cc_so_name = "__init__",
     cc_srcs = [
@@ -180,6 +182,7 @@ drake_py_unittest(
         "//bindings/pydrake/multibody:plant_py",
         "//bindings/pydrake/systems:analysis_py",
         "//bindings/pydrake/systems:lcm_py",
+        "//bindings/pydrake/systems:test_util_py",
     ],
 )
 

bindings/pydrake/geometry/geometry_py_visualizers.cc

@@ -3,11 +3,13 @@
 
 #include "drake/bindings/pydrake/common/default_scalars_pybind.h"
 #include "drake/bindings/pydrake/common/deprecation_pybind.h"
+#include "drake/bindings/pydrake/common/ref_cycle_pybind.h"
 #include "drake/bindings/pydrake/common/serialize_pybind.h"
 #include "drake/bindings/pydrake/common/type_pack.h"
 #include "drake/bindings/pydrake/common/type_safe_index_pybind.h"
 #include "drake/bindings/pydrake/documentation_pybind.h"
 #include "drake/bindings/pydrake/geometry/geometry_py.h"
+#include "drake/bindings/pydrake/systems/builder_life_support_pybind.h"
 #include "drake/geometry/drake_visualizer.h"
 #include "drake/geometry/meshcat.h"
 #include "drake/geometry/meshcat_animation.h"
@@ -53,10 +55,14 @@ void DefineDrakeVisualizer(py::module m, T) {
             py::arg("builder"), py::arg("scene_graph"),
             py::arg("lcm") = nullptr,
             py::arg("params") = DrakeVisualizerParams{},
-            // Keep alive, ownership: `return` keeps `builder` alive.
-            py::keep_alive<0, 1>(),
-            // Keep alive, reference: `builder` keeps `lcm` alive.
-            py::keep_alive<1, 3>(), py_rvp::reference,
+            // `return` and `builder` join ref cycle.
+            internal::ref_cycle<0, 1>(),
+            // Using builder_life_support_stash makes the builder temporarily
+            // immortal (uncollectible self cycle). This will be resolved by
+            // the Build() step. See BuilderLifeSupport for rationale.
+            internal::builder_life_support_stash<T, 1>(),
+            // Keep alive, reference: `return` keeps `lcm` alive.
+            py::keep_alive<0, 3>(), py_rvp::reference,
             cls_doc.AddToBuilder.doc_4args_builder_scene_graph_lcm_params)
         .def_static("AddToBuilder",
             py::overload_cast<systems::DiagramBuilder<T>*,
@@ -65,10 +71,14 @@ void DefineDrakeVisualizer(py::module m, T) {
             py::arg("builder"), py::arg("query_object_port"),
             py::arg("lcm") = nullptr,
             py::arg("params") = DrakeVisualizerParams{},
-            // Keep alive, ownership: `return` keeps `builder` alive.
-            py::keep_alive<0, 1>(),
-            // Keep alive, reference: `builder` keeps `lcm` alive.
-            py::keep_alive<1, 3>(), py_rvp::reference,
+            // `return` and `builder` join ref cycle.
+            internal::ref_cycle<0, 1>(),
+            // Using builder_life_support_stash makes the builder temporarily
+            // immortal (uncollectible self cycle). This will be resolved by
+            // the Build() step. See BuilderLifeSupport for rationale.
+            internal::builder_life_support_stash<T, 1>(),
+            // Keep alive, reference: `return` keeps `lcm` alive.
+            py::keep_alive<0, 3>(), py_rvp::reference,
             cls_doc.AddToBuilder.doc_4args_builder_query_object_port_lcm_params)
         .def_static("DispatchLoadMessage",
             &DrakeVisualizer<T>::DispatchLoadMessage, py::arg("scene_graph"),

bindings/pydrake/geometry/test/visualizers_test.py

@@ -1,8 +1,11 @@
 import pydrake.geometry as mut
 
 import copy
+import gc
+import itertools
 import unittest
 import urllib.request
+import weakref
 
 import numpy as np
 import umsgpack
@@ -15,6 +18,7 @@
 from pydrake.perception import PointCloud
 from pydrake.systems.analysis import Simulator_
 from pydrake.systems.framework import DiagramBuilder_, InputPort_
+from pydrake.systems.test.test_util import call_build_from_cpp
 
 # TODO(mwoehlke-kitware): Remove this when Jammy's python3-u-msgpack has been
 # updated to 2.5.2 or later.
@@ -30,7 +34,6 @@ def test_drake_visualizer(self, T):
         SceneGraph = mut.SceneGraph_[T]
         DiagramBuilder = DiagramBuilder_[T]
         Simulator = Simulator_[T]
-        lcm = DrakeLcm()
         role = mut.Role.kIllustration
         params = mut.DrakeVisualizerParams(
             publish_period=0.1, role=mut.Role.kIllustration,
@@ -41,47 +44,96 @@ def test_drake_visualizer(self, T):
         copy.copy(params)
 
         # Add some subscribers to detect message broadcast.
-        load_channel = "DRAKE_VIEWER_LOAD_ROBOT"
-        draw_channel = "DRAKE_VIEWER_DRAW"
-        load_subscriber = Subscriber(
-            lcm, load_channel, lcmt_viewer_load_robot)
-        draw_subscriber = Subscriber(
-            lcm, draw_channel, lcmt_viewer_draw)
-
-        # There are three ways to configure DrakeVisualizer.
+        def add_subscribers(lcm):
+            load_channel = "DRAKE_VIEWER_LOAD_ROBOT"
+            draw_channel = "DRAKE_VIEWER_DRAW"
+            load_subscriber = Subscriber(
+                lcm, load_channel, lcmt_viewer_load_robot)
+            draw_subscriber = Subscriber(
+                lcm, draw_channel, lcmt_viewer_draw)
+            return load_subscriber, draw_subscriber
+
+        # There are three ways to place DrakeVisualizer into a Diagram. We'll
+        # declare a helper functions for each of the ways. All of the helpers
+        # return a weak reference to the DrakeLcm object that they create, so
+        # that the test logic can verify correct object lifetimes.
         def by_hand(builder, scene_graph, params):
+            lcm = DrakeLcm()
             visualizer = builder.AddSystem(
                 mut.DrakeVisualizer_[T](lcm=lcm, params=params))
             builder.Connect(scene_graph.get_query_output_port(),
                             visualizer.query_object_input_port())
+            return weakref.finalize(lcm, lambda: None)
 
         def auto_connect_to_system(builder, scene_graph, params):
+            lcm = DrakeLcm()
             mut.DrakeVisualizer_[T].AddToBuilder(builder=builder,
                                                  scene_graph=scene_graph,
                                                  lcm=lcm, params=params)
+            return weakref.finalize(lcm, lambda: None)
 
         def auto_connect_to_port(builder, scene_graph, params):
+            lcm = DrakeLcm()
             mut.DrakeVisualizer_[T].AddToBuilder(
                 builder=builder,
                 query_object_port=scene_graph.get_query_output_port(),
                 lcm=lcm, params=params)
+            return weakref.finalize(lcm, lambda: None)
+
+        # After adding the visualizer, there are two ways to build the
+        # diagram.
+
+        def python_builder_build(builder):
+            # The pydrake binding ensures object lifetimes properly.
+            return builder.Build()
+
+        def cxx_builder_build(builder):
+            # Calling build on a python-populated diagram builder from c++ is
+            # dodgy at best, but don't worry -- all will be well if something
+            # retains ownership of the diagram.
+            return call_build_from_cpp(builder)
 
-        for func in [by_hand, auto_connect_to_system, auto_connect_to_port]:
+        for add_function, build_function in itertools.product(
+                    [by_hand, auto_connect_to_system, auto_connect_to_port],
+                    [python_builder_build, cxx_builder_build]):
             # Build the diagram.
             builder = DiagramBuilder()
             scene_graph = builder.AddSystem(SceneGraph())
-            func(builder, scene_graph, params)
+            lcm_spy = add_function(builder, scene_graph, params)
+            diagram = build_function(builder)
+
+            # Remove extraneous references, and collect garbage for good
+            # measure.
+            del builder
+            del scene_graph
+            gc.collect()
+
+            # Check the lcm object is still alive, recover a reference, and set
+            # receivers.
+            assert lcm_spy.alive
+            lcm = lcm_spy.peek()[0]
+            load_subscriber, draw_subscriber = add_subscribers(lcm)
 
             # Simulate to t = 0 to send initial load and draw messages.
-            diagram = builder.Build()
             Simulator(diagram).AdvanceTo(0)
             lcm.HandleSubscriptions(0)
             self.assertEqual(load_subscriber.count, 1)
             self.assertEqual(draw_subscriber.count, 1)
             load_subscriber.clear()
             draw_subscriber.clear()
 
+            # The diagram is mortal.
+            diagram_spy = weakref.finalize(diagram, lambda: None)
+            del diagram
+            del lcm
+            gc.collect()
+            self.assertFalse(diagram_spy.alive)
+            self.assertFalse(lcm_spy.alive)
+
         # Ad hoc broadcasting.
+        # Recreate some items deleted above.
+        lcm = DrakeLcm()
+        load_subscriber, draw_subscriber = add_subscribers(lcm)
         scene_graph = SceneGraph()
 
         mut.DrakeVisualizer_[T].DispatchLoadMessage(

bindings/pydrake/manipulation/BUILD.bazel

@@ -25,7 +25,9 @@ drake_pybind_library(
     cc_deps = [
         "//bindings/pydrake:documentation_pybind",
         "//bindings/pydrake/common:eigen_pybind",
+        "//bindings/pydrake/common:ref_cycle_pybind",
         "//bindings/pydrake/common:serialize_pybind",
+        "//bindings/pydrake/systems:builder_life_support_pybind",
     ],
     cc_so_name = "__init__",
     cc_srcs = [
@@ -56,6 +58,7 @@ drake_py_unittest(
         ":manipulation",
         "//bindings/pydrake/multibody",
         "//bindings/pydrake/systems",
+        "//bindings/pydrake/systems:test_util_py",
     ],
 )
 

bindings/pydrake/manipulation/manipulation_py_kuka_iiwa.cc

@@ -1,7 +1,9 @@
+#include "drake/bindings/pydrake/common/ref_cycle_pybind.h"
 #include "drake/bindings/pydrake/common/serialize_pybind.h"
 #include "drake/bindings/pydrake/documentation_pybind.h"
 #include "drake/bindings/pydrake/manipulation/manipulation_py.h"
 #include "drake/bindings/pydrake/pydrake_pybind.h"
+#include "drake/bindings/pydrake/systems/builder_life_support_pybind.h"
 #include "drake/manipulation/kuka_iiwa/build_iiwa_control.h"
 #include "drake/manipulation/kuka_iiwa/iiwa_command_receiver.h"
 #include "drake/manipulation/kuka_iiwa/iiwa_command_sender.h"
@@ -179,10 +181,15 @@ void DefineManipulationKukaIiwa(py::module m) {
             py::arg("plant"), py::arg("iiwa_instance"),
             py::arg("controller_plant"), py::arg("ext_joint_filter_tau"),
             py::arg("desired_iiwa_kp_gains"), py::arg("control_mode"),
-            // Keep alive, ownership: `return` keeps `builder` alive.
-            py::keep_alive<0, 1>(),
-            // Keep alive, reference: `builder` keeps `controller_plant` alive.
-            py::keep_alive<1, 4>(), py_rvp::reference,
+            // Using builder_life_support_stash makes the
+            // builder temporarily immortal (uncollectible self cycle). This
+            // will be resolved by the Build() step. See BuilderLifeSupport
+            // for rationale.
+            internal::builder_life_support_stash<double, 1>(),
+            // `return` and `builder` join ref cycle.
+            internal::ref_cycle<0, 1>(),
+            // Keep alive, reference: `return` keeps `controller_plant` alive.
+            py::keep_alive<0, 4>(), py_rvp::reference,
             cls_doc.AddToBuilder.doc);
   }
 
@@ -197,7 +204,16 @@ void DefineManipulationKukaIiwa(py::module m) {
         py::arg("builder"), py::arg("ext_joint_filter_tau") = 0.01,
         py::arg("desired_iiwa_kp_gains") = std::nullopt,
         py::arg("control_mode") = IiwaControlMode::kPositionAndTorque,
-        // Keep alive, reference: `builder` keeps `controller_plant` alive.
+        // Using builder_life_support_stash makes the
+        // builder temporarily immortal (uncollectible self cycle). This
+        // will be resolved by the Build() step. See BuilderLifeSupport
+        // for rationale.
+        internal::builder_life_support_stash<double, 5>(),
+        // Keep alive, reference: `builder` keeps `controller_plant` alive.  It
+        // would be preferable to attach the keep-alive to the SimIiwaDriver
+        // diagram/system, but it does not appear in the function signature.
+        // Use the builder as an adequate lifetime proxy, since it will be kept
+        // alive via lifetime management associated with Build() calls.
         py::keep_alive<5, 3>(), doc.BuildIiwaControl.doc);
   }
 }