Allow for reseting the Quent context in StreamingEngine._reset.

python/cudf_polars/cudf_polars/engine/core.py

@@ -539,12 +539,6 @@ def _reset(
         existing_executor_options = self.config.get("executor_options", {})
         if not isinstance(existing_executor_options, dict):
             existing_executor_options = {}
-        existing_quent_context = existing_executor_options.get("quent_context")
-        if (
-            "quent_context" in executor_options
-            and executor_options["quent_context"] != existing_quent_context
-        ):
-            raise ValueError("quent_context cannot be changed during reset")
 
         _disallowed_exec = {"num_py_executors"} & executor_options.keys()
         if _disallowed_exec:
@@ -589,6 +583,42 @@ def __exit__(self, *_: object) -> None:
         """Exit the context manager, calling :meth:`shutdown`."""
         self.shutdown()
 
+    def _update_quent(
+        self,
+        new_quent_context: QuentConfig | None,
+        existing_quent_context: QuentConfig | None,
+        executor_options: dict[str, Any],
+    ) -> None:
+        """
+        Apply a ``quent_context`` change in-place.
+
+        Closes existing Quent runtimes, starts new ones for
+        ``new_quent_context``, and patches ``executor_options["quent_context"]``
+        and ``_quent_output_root``.  Subclasses implement :meth:`_quent_close`
+        and :meth:`_quent_open` for the backend-specific work.
+        """
+        if new_quent_context != existing_quent_context:
+            if existing_quent_context is not None:
+                self._quent_close()
+            if new_quent_context is not None:
+                new_quent_context = self._quent_open(new_quent_context)
+        self._quent_output_root = (
+            new_quent_context.run_root if new_quent_context is not None else None
+        )
+        executor_options["quent_context"] = new_quent_context
+
+    def _quent_close(self) -> None:
+        """Close all Quent runtimes for this engine (workers then controller)."""
+        raise NotImplementedError(  # pragma: no cover
+            f"{type(self).__name__} does not implement _quent_close"
+        )
+
+    def _quent_open(self, quent_context: QuentConfig) -> QuentConfig:
+        """Start fresh Quent runtimes and return the effective configuration."""
+        raise NotImplementedError(  # pragma: no cover
+            f"{type(self).__name__} does not implement _quent_open"
+        )
+
     def _run(self, func: Callable[..., T], *args: Any, **kwargs: Any) -> list[T]:
         """
         Execute a function on all ranks.