chore: import upstream snapshot with attribution

python/ray/llm/_internal/batch/observability/__init__.py

@@ -0,0 +1,23 @@
+from ray.llm._internal.batch.observability.logging.setup import (
+    setup_logging,
+)
+from ray.llm._internal.common.observability.logging_utils import (
+    disable_datasets_logger,
+    disable_vllm_custom_ops_logger_on_cpu_nodes,
+)
+from ray.llm._internal.common.observability.telemetry_utils import Once
+
+_setup_observability_once = Once()
+
+
+def _setup_observability():
+    setup_logging()
+    disable_datasets_logger()
+    disable_vllm_custom_ops_logger_on_cpu_nodes()
+
+
+def setup_observability():
+    _setup_observability_once.do_once(_setup_observability)
+
+
+__all__ = ["setup_observability"]

python/ray/llm/_internal/batch/observability/logging/__init__.py

@@ -0,0 +1,40 @@
+import logging
+from typing import Optional
+
+from ray._common.filters import CoreContextFilter
+
+
+def _setup_logger(logger_name: str) -> None:
+    """Setup logger given the logger name.
+
+    This function is idempotent and won't set up the same logger multiple times. It will
+    also skip the setup if Data logger is already setup and has handlers.
+
+    Args:
+        logger_name: logger name used to get the logger.
+    """
+    logger = logging.getLogger(logger_name)
+    llm_logger = logging.getLogger("ray.data")
+
+    # Skip setup if the logger already has handlers setup or if the parent (Data
+    # logger) has handlers.
+    if logger.handlers or llm_logger.handlers:
+        return
+
+    # Set up stream handler, which logs to console as plaintext.
+    stream_handler = logging.StreamHandler()
+    stream_handler.addFilter(CoreContextFilter())
+    logger.addHandler(stream_handler)
+    logger.setLevel(logging.INFO)
+    logger.propagate = False
+
+
+def get_logger(name: Optional[str] = None):
+    """Get a structured logger inherited from the Ray Data logger.
+
+    Loggers by default are logging to stdout, and are expected to be scraped by an
+    external process.
+    """
+    logger_name = f"ray.data.{name}"
+    _setup_logger(logger_name)
+    return logging.getLogger(logger_name)

python/ray/llm/_internal/batch/observability/logging/setup.py

@@ -0,0 +1,26 @@
+import logging
+
+from ray._common.filters import CoreContextFilter
+from ray._common.formatters import JSONFormatter
+
+
+def _configure_stdlib_logging():
+    """Configures stdlib root logger to make sure stdlib loggers (created as
+    `logging.getLogger(...)`) are using Ray's `JSONFormatter` with Core and Serve
+     context filters.
+    """
+
+    handler = logging.StreamHandler()
+    handler.addFilter(CoreContextFilter())
+    handler.setFormatter(JSONFormatter())
+
+    root_logger = logging.getLogger()
+    # NOTE: It's crucial we reset all the handlers of the root logger,
+    #       to make sure that logs aren't emitted twice
+    root_logger.handlers = []
+    root_logger.addHandler(handler)
+    root_logger.setLevel(logging.INFO)
+
+
+def setup_logging():
+    _configure_stdlib_logging()

python/ray/llm/_internal/batch/observability/usage_telemetry/usage.py

@@ -0,0 +1,162 @@
+from enum import Enum
+from typing import Callable, Dict, Tuple, Union
+
+import ray
+from ray._common.constants import HEAD_NODE_RESOURCE_NAME
+from ray._common.usage.usage_lib import record_extra_usage_tag
+from ray.llm._internal.batch.observability.logging import get_logger
+from ray.llm._internal.common.base_pydantic import BaseModelExtended
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+LLM_BATCH_TELEMETRY_NAMESPACE = "llm_batch_telemetry"
+LLM_BATCH_TELEMETRY_ACTOR_NAME = "llm_batch_telemetry"
+
+logger = get_logger(__name__)
+
+
+class BatchModelTelemetry(BaseModelExtended):
+    # Dedup identity only; never recorded as a tag value. A hash of model_source
+    # so distinct models that share an architecture stay separate in the dedup
+    # key while the cleartext model name never reaches the head-node actor.
+    model_id_hash: str = ""
+    processor_config_name: str = ""
+    model_architecture: str = ""
+    batch_size: int = 0
+    accelerator_type: str = ""
+    concurrency: Union[int, Tuple[int, int]] = 0
+    task_type: str = ""
+
+    # For the parallel size, 0 means not supported.
+    pipeline_parallel_size: int = 0
+    tensor_parallel_size: int = 0
+    data_parallel_size: int = 0
+
+
+class BatchTelemetryTags(str, Enum):
+    """Telemetry tags for RayLLM Batch."""
+
+    LLM_BATCH_PROCESSOR_CONFIG_NAME = "LLM_BATCH_PROCESSOR_CONFIG_NAME"
+    LLM_BATCH_MODEL_ARCHITECTURE = "LLM_BATCH_MODEL_ARCHITECTURE"
+    LLM_BATCH_SIZE = "LLM_BATCH_SIZE"
+    LLM_BATCH_ACCELERATOR_TYPE = "LLM_BATCH_ACCELERATOR_TYPE"
+    LLM_BATCH_CONCURRENCY = "LLM_BATCH_CONCURRENCY"
+    LLM_BATCH_TASK_TYPE = "LLM_BATCH_TASK_TYPE"
+    LLM_BATCH_PIPELINE_PARALLEL_SIZE = "LLM_BATCH_PIPELINE_PARALLEL_SIZE"
+    LLM_BATCH_TENSOR_PARALLEL_SIZE = "LLM_BATCH_TENSOR_PARALLEL_SIZE"
+    LLM_BATCH_DATA_PARALLEL_SIZE = "LLM_BATCH_DATA_PARALLEL_SIZE"
+
+
+@ray.remote(
+    name=LLM_BATCH_TELEMETRY_ACTOR_NAME,
+    namespace=LLM_BATCH_TELEMETRY_NAMESPACE,
+    num_cpus=0,
+    lifetime="detached",
+)
+class _TelemetryAgent:
+    """Named Actor to keep the state of all deployed models and record telemetry."""
+
+    def __init__(self):
+        # Keyed by full telemetry identity (incl. model_id_hash) so repeated
+        # identical processor builds overwrite while distinct models/configs
+        # remain separate.
+        self._tracking_telemetries: Dict[str, BatchModelTelemetry] = {}
+        self._record_tag_func = record_extra_usage_tag
+
+    def _update_record_tag_func(self, record_tag_func: Callable) -> None:
+        self._record_tag_func = record_tag_func
+
+    def _reset(self) -> None:
+        """Only used in tests to clear accumulated telemetries."""
+        self._tracking_telemetries = {}
+
+    def generate_report(self) -> Dict[str, str]:
+        return {
+            BatchTelemetryTags.LLM_BATCH_PROCESSOR_CONFIG_NAME: ",".join(
+                [t.processor_config_name for t in self._tracking_telemetries.values()]
+            ),
+            BatchTelemetryTags.LLM_BATCH_MODEL_ARCHITECTURE: ",".join(
+                [t.model_architecture for t in self._tracking_telemetries.values()]
+            ),
+            BatchTelemetryTags.LLM_BATCH_SIZE: ",".join(
+                [str(t.batch_size) for t in self._tracking_telemetries.values()]
+            ),
+            BatchTelemetryTags.LLM_BATCH_ACCELERATOR_TYPE: ",".join(
+                [t.accelerator_type for t in self._tracking_telemetries.values()]
+            ),
+            BatchTelemetryTags.LLM_BATCH_CONCURRENCY: ",".join(
+                [str(t.concurrency) for t in self._tracking_telemetries.values()]
+            ),
+            BatchTelemetryTags.LLM_BATCH_TASK_TYPE: ",".join(
+                [t.task_type for t in self._tracking_telemetries.values()]
+            ),
+            BatchTelemetryTags.LLM_BATCH_PIPELINE_PARALLEL_SIZE: ",".join(
+                [
+                    str(t.pipeline_parallel_size)
+                    for t in self._tracking_telemetries.values()
+                ]
+            ),
+            BatchTelemetryTags.LLM_BATCH_TENSOR_PARALLEL_SIZE: ",".join(
+                [
+                    str(t.tensor_parallel_size)
+                    for t in self._tracking_telemetries.values()
+                ]
+            ),
+            BatchTelemetryTags.LLM_BATCH_DATA_PARALLEL_SIZE: ",".join(
+                [str(t.data_parallel_size) for t in self._tracking_telemetries.values()]
+            ),
+        }
+
+    def record(self, telemetry: BatchModelTelemetry) -> None:
+        """Upsert by identity and record telemetries."""
+        from ray._common.usage.usage_lib import TagKey
+
+        self._tracking_telemetries[telemetry.model_dump_json()] = telemetry
+        for key, value in self.generate_report().items():
+            try:
+                self._record_tag_func(TagKey.Value(key), value)
+            except ValueError:
+                # Tag not in the installed usage proto; skip rather than fail.
+                continue
+
+
+class TelemetryAgent:
+    """Wrapper around the telemetry agent that calls the remote method until
+    push_telemetry_report is called."""
+
+    def __init__(self):
+        # Creation must never break processor construction, so swallow failures
+        # (e.g. Ray not initialized, transient GCS issues) and degrade to a no-op.
+        try:
+            # get_if_exists makes creation atomic across concurrent drivers.
+            self.remote_telemetry_agent = _TelemetryAgent.options(
+                name=LLM_BATCH_TELEMETRY_ACTOR_NAME,
+                namespace=LLM_BATCH_TELEMETRY_NAMESPACE,
+                get_if_exists=True,
+                # Ensure the actor is created on the head node.
+                resources={HEAD_NODE_RESOURCE_NAME: 0.001},
+                # Ensure the actor is not scheduled with the existing placement group.
+                scheduling_strategy=PlacementGroupSchedulingStrategy(
+                    placement_group=None
+                ),
+            ).remote()
+        except Exception:
+            self.remote_telemetry_agent = None
+            logger.exception("Failed to initialize LLM batch telemetry agent")
+
+    def _update_record_tag_func(self, record_tag_func: Callable):
+        if self.remote_telemetry_agent is not None:
+            self.remote_telemetry_agent._update_record_tag_func.remote(record_tag_func)
+
+    def push_telemetry_report(self, telemetry: BatchModelTelemetry):
+        # Telemetry must never break processor construction.
+        if self.remote_telemetry_agent is None:
+            return
+        try:
+            ray.get(self.remote_telemetry_agent.record.remote(telemetry))
+        except Exception:
+            logger.exception("Failed to push LLM batch telemetry")
+
+
+def get_or_create_telemetry_agent() -> TelemetryAgent:
+    """Helper to get or create the telemetry agent."""
+    return TelemetryAgent()