Merge pull request #22155 from philip-stoev/mzcompose_skip_sanity

mzcompose: Allow tests to opt out of the mz_sanity_restart check

Author: philip-stoev

misc/python/materialize/mzcompose/composition.py

@@ -701,13 +701,13 @@ def validate_sources_sinks_clusters(self) -> str | None:
     def sanity_restart_mz(self) -> None:
         """Restart Materialized if it is part of the composition to find
         problems with persisted objects, functions as a sanity check."""
-        # Exclude environmentd image, which is used in cloud-canary, and doesn't start up fully without clusterd images
         if (
             "materialized" in self.compose["services"]
-            and "materialize/environmentd"
-            not in self.compose["services"]["materialized"]["image"]
+            and "sanity_restart" in self.compose["services"]["materialized"]["labels"]
         ):
-            ui.header("Sanity Check: Restart Mz, verify source/sink/replica health")
+            ui.header(
+                "Sanity Restart: Restart Mz and verify source/sink/replica health"
+            )
             self.kill("materialized")
             # TODO(def-): Better way to detect when kill has finished
             time.sleep(3)
@@ -724,6 +724,10 @@ def sanity_restart_mz(self) -> None:
                 # Sources and cluster replicas need a few seconds to start up
                 print(f"Retrying ({i+1}/{NUM_RETRIES})...")
                 time.sleep(1)
+        else:
+            ui.header(
+                "Sanity Restart skipped because Mz not in services or `sanity_restart` label not set"
+            )
 
     def down(
         self,

misc/python/materialize/mzcompose/service.py

@@ -159,6 +159,9 @@ class ServiceConfig(TypedDict, total=False):
     restart: str
     """Restart policy."""
 
+    labels: list[str]
+    """Container labels."""
+
 
 class Service:
     """A Docker Compose service in a `Composition`.

misc/python/materialize/mzcompose/services/materialized.py

@@ -48,6 +48,7 @@ def __init__(
         system_parameter_defaults: dict[str, str] | None = None,
         additional_system_parameter_defaults: dict[str, str] | None = None,
         soft_assertions: bool = True,
+        sanity_restart: bool = True,
     ) -> None:
         depends_graph: dict[str, ServiceDependency] = {
             s: {"condition": "service_started"} for s in depends_on
@@ -164,6 +165,9 @@ def __init__(
         if memory:
             config["deploy"] = {"resources": {"limits": {"memory": memory}}}
 
+        if sanity_restart:
+            config.setdefault("labels", []).append("sanity_restart")
+
         volumes = []
         if use_default_volumes:
             volumes += DEFAULT_MZ_VOLUMES

test/cloud-canary/mzcompose.py

@@ -54,6 +54,8 @@
             "--orchestrator-process-secrets-directory=/mzdata/secrets",
             "--orchestrator-process-scratch-directory=/scratch",
         ],
+        # We can not restart this container at will, as it does not have clusterd
+        sanity_restart=False,
     ),
     Testdrive(),  # Overriden below
     Mz(