chore: import upstream snapshot with attribution

python/ray/train/v2/_internal/metrics/worker.py

@@ -0,0 +1,64 @@
+from typing import Dict
+
+from ray.train.v2._internal.metrics.base import (
+    RUN_ID_TAG_KEY,
+    RUN_NAME_TAG_KEY,
+    TimeMetric,
+)
+
+WORKER_WORLD_RANK_TAG_KEY = "ray_train_worker_world_rank"
+WORKER_ACTOR_ID_TAG_KEY = "ray_train_worker_actor_id"
+
+
+class WorkerMetrics:
+    """Factory for creating worker-specific metrics.
+
+    This class defines all metrics used to track the state and performance of the
+    training workers. Each metric is defined with its name, type, default value,
+    description, and required tags.
+    """
+
+    # ===== Metric Names =====
+    REPORT_TOTAL_BLOCKED_TIME_S = "train_report_total_blocked_time_s"
+    CHECKPOINT_SYNC_TOTAL_TIME_S = "train_checkpoint_sync_total_time_s"
+    CHECKPOINT_TRANSFER_TOTAL_TIME_S = "train_checkpoint_transfer_total_time_s"
+
+    @classmethod
+    def _create_time_metric(
+        cls, name: str, description: str, base_tags: Dict[str, str]
+    ) -> TimeMetric:
+        """Create a time-based metric."""
+        return TimeMetric(
+            name=name,
+            description=description,
+            base_tags=base_tags,
+        )
+
+    @classmethod
+    def get_worker_metrics(
+        cls, run_name: str, run_id: str, world_rank: int, worker_actor_id: str
+    ) -> Dict[str, TimeMetric]:
+        """Get all worker metrics."""
+        base_tags = {
+            RUN_NAME_TAG_KEY: run_name,
+            RUN_ID_TAG_KEY: run_id,
+            WORKER_WORLD_RANK_TAG_KEY: str(world_rank),
+            WORKER_ACTOR_ID_TAG_KEY: worker_actor_id,
+        }
+        return {
+            cls.REPORT_TOTAL_BLOCKED_TIME_S: cls._create_time_metric(
+                cls.REPORT_TOTAL_BLOCKED_TIME_S,
+                "Cumulative time in seconds to report a checkpoint to the storage.",
+                base_tags,
+            ),
+            cls.CHECKPOINT_SYNC_TOTAL_TIME_S: cls._create_time_metric(
+                cls.CHECKPOINT_SYNC_TOTAL_TIME_S,
+                "Cumulative time in seconds spent synchronizing the checkpoint directory name across all ranks.",
+                base_tags,
+            ),
+            cls.CHECKPOINT_TRANSFER_TOTAL_TIME_S: cls._create_time_metric(
+                cls.CHECKPOINT_TRANSFER_TOTAL_TIME_S,
+                "Cumulative time in seconds spent transferring checkpoint files to storage.",
+                base_tags,
+            ),
+        }