Repository: ray-project/ray
File: python/ray/data/_internal/execution/streaming_executor.py (Python)

Code before:
            "Budget (Object Store Memory) per operator",
            tag_keys=("dataset", "operator"),
        )
        self._max_bytes_to_read_gauge: Gauge = Gauge(
            "data_max_bytes_to_read",
            description="Maximum bytes to read from streaming generator buffer.",
            tag_keys=("dataset", "operator"),
        )

        # Cluster autoscaler utilization gauges
        self._cluster_cpu_utilization_gauge: Gauge = Gauge(
            "data_cluster_cpu_utilization",
            description="Cluster utilization % (CPU)",
            tag_keys=("dataset",),
        )
        self._cluster_gpu_utilization_gauge: Gauge = Gauge(
            "data_cluster_gpu_utilization",
            description="Cluster utilization % (GPU)",
            tag_keys=("dataset",),
        )
        self._cluster_object_store_memory_utilization_gauge: Gauge = Gauge(
            "data_cluster_object_store_memory_utilization",
            description="Cluster utilization % (Object Store Memory)",
            tag_keys=("dataset",),
        )

    def execute(
        self, dag: PhysicalOperator, initial_stats: Optional[DatasetStats] = None
    ) -> OutputIterator:
        """Executes the DAG using a streaming execution strategy.

        We take an event-loop approach to scheduling. We block on the next scheduling
        event using `ray.wait`, updating operator state and dispatching new tasks.
        """

        self._initial_stats = initial_stats
        self._start_time = time.perf_counter()

        if not isinstance(dag, InputDataBuffer):
            if self._data_context.print_on_execution_start:
                message = f"Starting execution of Dataset {self._dataset_id}."
                log_path = get_log_directory()
                if log_path is not None:
                    message += f" Full logs are in {log_path}"
                logger.info(message)
                logger.info(
                    f"Execution plan of Dataset {self._dataset_id}: {dag.dag_str}"
                )

            logger.debug("Execution config: %s", self._options)


Review comment (refactor):
These are metrics that are specific to an autoscaling implementation detail. Rather than placing this in the streaming executor (core scheduling logic), could you move this to `RollingLogicalUtilizationGauge`? I think we can add an optional `execution_id` parameter to the constructor, define the gauges in the constructor, and if the `execution_id` isn't None, we can update update the metrics

Code after:
            "data_memory_budget",
            "Budget (Memory) per operator",
            tag_keys=("dataset", "operator"),
        )
        self._osm_budget_gauge: Gauge = Gauge(
            "data_object_store_memory_budget",
            "Budget (Object Store Memory) per operator",
            tag_keys=("dataset", "operator"),
        )
        self._max_bytes_to_read_gauge: Gauge = Gauge(
            "data_max_bytes_to_read",
            description="Maximum bytes to read from streaming generator buffer.",
            tag_keys=("dataset", "operator"),
        )

    def execute(
        self, dag: PhysicalOperator, initial_stats: Optional[DatasetStats] = None
    ) -> OutputIterator:
        """Executes the DAG using a streaming execution strategy.

        We take an event-loop approach to scheduling. We block on the next scheduling
        event using `ray.wait`, updating operator state and dispatching new tasks.
        """

        self._initial_stats = initial_stats
        self._start_time = time.perf_counter()

        if not isinstance(dag, InputDataBuffer):
            if self._data_context.print_on_execution_start:
                message = f"Starting execution of Dataset {self._dataset_id}."
                log_path = get_log_directory()
                if log_path is not None:
                    message += f" Full logs are in {log_path}"
                logger.info(message)
                logger.info(
                    f"Execution plan of Dataset {self._dataset_id}: {dag.dag_str}"
                )

            logger.debug("Execution config: %s", self._options)

        # Setup the streaming DAG topology and start the runner thread.
        self._topology = build_streaming_topology(dag, self._options)

        self._resource_manager = ResourceManager(
            self._topology,
            self._options,
            lambda: self._cluster_autoscaler.get_total_resources(),
            self._data_context,
        )

        # Setup progress manager