[utilization] pipeline to create clean db records (pytorch#145327)

yangw-dev · huydhn · pytorchmergebot · commit a9ed7bd78ea6 · 2025-01-29T23:48:50.000Z
upload_utilization_script to generate db-ready-insert records to s3 - generate two files: metadata and timeseries in ossci-utilization buckets - convert log record to db format ones - add unit test job for tools/stats/ Related Prs: setup composite action for data pipeline: pytorch#145310 add permission for composite action to access S3 bucket: meta-pytorch/pytorch-gha-infra#595 add insert logic in s3 replicator: pytorch/test-infra#6217 Pull Request resolved: pytorch#145327 Approved by: https://github.com/huydhn Co-authored-by: Huy Do <huydhn@gmail.com>
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
@@ -1,4 +1,5 @@
 name: Lint
+# Workflow that runs lint checks and also unittests for tools, and scripts.
 
 on:
   pull_request:
@@ -207,6 +208,7 @@ jobs:
         conda activate "${CONDA_ENV}"
 
         # Test tools
+        PYTHONPATH=$(pwd) pytest tools/stats
         PYTHONPATH=$(pwd) pytest tools/test -o "python_files=test*.py"
         PYTHONPATH=$(pwd) pytest .github/scripts -o "python_files=test*.py"
 
diff --git a/.lintrunner.toml b/.lintrunner.toml
@@ -161,7 +161,8 @@ init_command = [
     'rich==10.9.0',
     'pyyaml==6.0.1',
     'optree==0.13.0',
-    'dataclasses_json==0.6.7'
+    'dataclasses_json==0.6.7',
+    'pandas==2.2.3',
 ]
 
 [[linter]]
diff --git a/tools/stats/monitor.py b/tools/stats/monitor.py
@@ -190,6 +190,7 @@ def __init__(
             job_name=_job_name,
             workflow_id=_workflow_run_id,
             workflow_name=_workflow_name,
+            start_at=datetime.datetime.now().timestamp(),
         )
         self._data_collect_interval = data_collect_interval
         self._has_pynvml = pynvml_enabled
@@ -257,7 +258,11 @@ def _output_data(self) -> None:
 
         while not self.exit_event.is_set():
             collecting_start_time = time.time()
-            stats = UtilizationRecord()
+            stats = UtilizationRecord(
+                level="record",
+                timestamp=datetime.datetime.now().timestamp(),
+            )
+
             try:
                 data_list, error_list = self.shared_resource.get_and_reset()
                 if self._debug_mode:
@@ -275,8 +280,6 @@ def _output_data(self) -> None:
                 if not data_list:
                     # pass since no data is collected
                     continue
-                stats.level = "record"
-                stats.timestamp = datetime.datetime.now().timestamp()
 
                 cpu_stats = self._generate_stats(
                     [data.cpu_percent for data in data_list]
diff --git a/tools/stats/upload_stats_lib.py b/tools/stats/upload_stats_lib.py
@@ -23,6 +23,9 @@ def get_s3_resource() -> Any:
     return boto3.resource("s3")
 
 
+GHA_ARTIFACTS_BUCKET = "gha-artifacts"
+
+
 # NB: In CI, a flaky test is usually retried 3 times, then the test file would be rerun
 # 2 more times
 MAX_RETRY_IN_NON_DISABLED_MODE = 3 * 3
@@ -84,16 +87,22 @@ def _download_artifact(
 
 
 def download_s3_artifacts(
-    prefix: str, workflow_run_id: int, workflow_run_attempt: int
+    prefix: str,
+    workflow_run_id: int,
+    workflow_run_attempt: int,
+    job_id: Optional[int] = None,
 ) -> list[Path]:
-    bucket = get_s3_resource().Bucket("gha-artifacts")
+    bucket = get_s3_resource().Bucket(GHA_ARTIFACTS_BUCKET)
     objs = bucket.objects.filter(
         Prefix=f"pytorch/pytorch/{workflow_run_id}/{workflow_run_attempt}/artifact/{prefix}"
     )
-
     found_one = False
     paths = []
     for obj in objs:
+        object_name = Path(obj.key).name
+        # target an artifact for a specific job_id if provided, otherwise skip the download.
+        if job_id is not None and str(job_id) not in object_name:
+            continue
         found_one = True
         p = Path(Path(obj.key).name)
         print(f"Downloading {p}")
diff --git a/tools/stats/upload_utilization_stats/test_upload_utilization_stats.py b/tools/stats/upload_utilization_stats/test_upload_utilization_stats.py
@@ -0,0 +1,195 @@
+import os
+import sys
+import unittest
+from collections import Counter
+from datetime import datetime, timedelta
+
+
+# adding sys.path makes the monitor script able to import path tools.stats.utilization_stats_lib
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", ".."))
+from tools.stats.upload_utilization_stats.upload_utilization_stats import (
+    SegmentGenerator,
+)
+from tools.stats.utilization_stats_lib import OssCiSegmentV1, UtilizationRecord
+
+
+# datetimes from January 1, 2022 12:00:00
+TEST_DT_BASE = datetime(2022, 1, 1, 12, 0, 0)
+TEST_DT_PLUS_5S = TEST_DT_BASE + timedelta(seconds=5)
+TEST_DT_PLUS_10S = TEST_DT_BASE + timedelta(seconds=10)
+TEST_DT_PLUS_15S = TEST_DT_BASE + timedelta(seconds=15)
+TEST_DT_PLUS_30S = TEST_DT_BASE + timedelta(seconds=30)
+TEST_DT_PLUS_40S = TEST_DT_BASE + timedelta(seconds=40)
+
+# timestamps from January 1, 2022 12:00:00
+TEST_TS_BASE = TEST_DT_BASE.timestamp()
+TEST_TS_PLUS_5S = TEST_DT_PLUS_5S.timestamp()
+TEST_TS_PLUS_10S = TEST_DT_PLUS_10S.timestamp()
+TEST_TS_PLUS_15S = TEST_DT_PLUS_15S.timestamp()
+TEST_TS_PLUS_30S = TEST_DT_PLUS_30S.timestamp()
+TEST_TS_PLUS_40S = TEST_DT_PLUS_40S.timestamp()
+
+# test cmd names
+PYTEST1_NAME = "python test1.py"
+PYTEST2_NAME = "python test2.py"
+PYPIP_INSTALL_NAME = "python pip install install1"
+
+
+class TestSegmentGenerator(unittest.TestCase):
+    def test_generate_empty_records(self) -> None:
+        records: list[UtilizationRecord] = []
+
+        # execute
+        generator = SegmentGenerator()
+        segments = generator.generate(records)
+
+        # assert
+        self.assertEqual(segments, [])
+
+    def test_generate_single_record(self) -> None:
+        record = UtilizationRecord(
+            timestamp=TEST_TS_BASE, cmd_names=[PYTEST1_NAME], level="PYTHON_CMD"
+        )
+        records = [record]
+
+        # execute
+        generator = SegmentGenerator()
+        segments = generator.generate(records)
+
+        # assert
+        self.assertEqual(len(segments), 1)
+
+    def test_generate_single_record_with_multiple_cmds(self) -> None:
+        record = UtilizationRecord(
+            timestamp=TEST_TS_BASE,
+            cmd_names=[PYTEST1_NAME, PYPIP_INSTALL_NAME],
+            level="PYTHON_CMD",
+        )
+        records = [record]
+
+        # execute
+        generator = SegmentGenerator()
+        segments = generator.generate(records)
+
+        # assert
+        self.assertEqual(len(segments), 2)
+
+    def test_generate_multiple_records(self) -> None:
+        records = get_base_test_records()
+
+        # execute
+        generator = SegmentGenerator()
+        segments = generator.generate(records)
+
+        # assert
+        self.assertEqual(len(segments), 2)
+        self.validate_segment(segments[0], PYTEST1_NAME, TEST_TS_BASE, TEST_TS_PLUS_30S)
+        self.validate_segment(
+            segments[1], PYPIP_INSTALL_NAME, TEST_TS_PLUS_10S, TEST_TS_PLUS_15S
+        )
+
+    def test_generate_cmd_interval_larger_than_default_threshold_setting(self) -> None:
+        records = get_base_test_records()
+
+        # record has more than 1 minute gap than last default record
+        test_gap_dt1 = TEST_DT_PLUS_30S + timedelta(seconds=80)
+        test_gap_dt2 = TEST_DT_PLUS_30S + timedelta(seconds=85)
+        record_gap_1 = UtilizationRecord(
+            timestamp=test_gap_dt1.timestamp(),
+            cmd_names=[PYTEST1_NAME],
+            level="PYTHON_CMD",
+        )
+        record_gap_2 = UtilizationRecord(
+            timestamp=test_gap_dt2.timestamp(),
+            cmd_names=[PYTEST1_NAME],
+            level="PYTHON_CMD",
+        )
+        records += [record_gap_1, record_gap_2]
+
+        # execute
+        generator = SegmentGenerator()
+        segments = generator.generate(records)
+
+        # assert
+        counter = Counter(seg.name for seg in segments)
+        self.assertEqual(counter[PYTEST1_NAME], 2)
+        self.assertEqual(counter[PYPIP_INSTALL_NAME], 1)
+        self.assertEqual(len(segments), 3)
+
+        self.validate_segment(segments[0], PYTEST1_NAME, TEST_TS_BASE, TEST_TS_PLUS_30S)
+        self.validate_segment(
+            segments[1],
+            PYTEST1_NAME,
+            test_gap_dt1.timestamp(),
+            test_gap_dt2.timestamp(),
+        )
+        self.validate_segment(
+            segments[2], PYPIP_INSTALL_NAME, TEST_TS_PLUS_10S, TEST_TS_PLUS_15S
+        )
+
+    def test_generate_multiple_segments_with_customized_threshold(self) -> None:
+        # set threshold to consider as continuous segment to 10 seconds
+        test_threshold = 10
+
+        records = get_base_test_records()
+
+        # execute
+        generator = SegmentGenerator()
+        segments = generator.generate(records, test_threshold)
+
+        # assert
+        counter = Counter(seg.name for seg in segments)
+        self.assertEqual(counter[PYTEST1_NAME], 2)
+        self.assertEqual(counter[PYPIP_INSTALL_NAME], 1)
+        self.assertEqual(len(segments), 3)
+
+        self.validate_segment(segments[0], PYTEST1_NAME, TEST_TS_BASE, TEST_TS_PLUS_15S)
+        self.validate_segment(
+            segments[1], PYTEST1_NAME, TEST_TS_PLUS_30S, TEST_TS_PLUS_30S
+        )
+        self.validate_segment(
+            segments[2], PYPIP_INSTALL_NAME, TEST_TS_PLUS_10S, TEST_TS_PLUS_15S
+        )
+
+    def validate_segment(
+        self, segment: OssCiSegmentV1, name: str, start_at: float, end_at: float
+    ) -> None:
+        self.assertEqual(segment.name, name)
+        self.assertEqual(segment.start_at, start_at)
+        self.assertEqual(segment.end_at, end_at)
+
+
+def get_base_test_records() -> list[UtilizationRecord]:
+    record1 = UtilizationRecord(
+        timestamp=TEST_TS_BASE, cmd_names=[PYTEST1_NAME], level="PYTHON_CMD"
+    )
+    record2 = UtilizationRecord(
+        timestamp=TEST_TS_PLUS_5S,
+        cmd_names=[PYTEST1_NAME],
+        level="PYTHON_CMD",
+    )
+    record3 = UtilizationRecord(
+        timestamp=TEST_TS_PLUS_10S,
+        cmd_names=[PYTEST1_NAME, PYPIP_INSTALL_NAME],
+        level="PYTHON_CMD",
+    )
+    record4 = UtilizationRecord(
+        timestamp=TEST_TS_PLUS_15S,
+        cmd_names=[PYTEST1_NAME, PYPIP_INSTALL_NAME],
+        level="PYTHON_CMD",
+    )
+    record5 = UtilizationRecord(
+        timestamp=TEST_TS_PLUS_30S,
+        cmd_names=[PYTEST1_NAME],
+        level="PYTHON_CMD",
+    )
+    record6 = UtilizationRecord(
+        timestamp=TEST_TS_PLUS_40S,
+        cmd_names=[],
+        level="PYTHON_CMD",
+    )
+    return [record1, record2, record3, record4, record5, record6]
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tools/stats/upload_utilization_stats/upload_utilization_stats.py b/tools/stats/upload_utilization_stats/upload_utilization_stats.py
diff --git a/tools/stats/utilization_stats_lib.py b/tools/stats/utilization_stats_lib.py

Original file line number	Diff line number	Diff line change
`@@ -161,7 +161,8 @@ init_command = [`
`161`	`161`	`'rich==10.9.0',`
`162`	`162`	`'pyyaml==6.0.1',`
`163`	`163`	`'optree==0.13.0',`
`164`		`- 'dataclasses_json==0.6.7'`
	`164`	`+ 'dataclasses_json==0.6.7',`
	`165`	`+ 'pandas==2.2.3',`
`165`	`166`	`]`
`166`	`167`
`167`	`168`	`[[linter]]`