yujunjun/datahub
2
0
Fork 0
mirror of https://github.com/datahub-project/datahub.git synced 2025-12-04 06:33:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(ingest): Track disk usage in report (#7812)

Browse Source
This commit is contained in:
Andrew Sikowitz 2023-04-13 17:43:25 -04:00 committed by GitHub
parent 4742e81726
commit ce795406b9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 22 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
metadata-ingestion/src/datahub/ingestion/run/pipeline.py
View File

@ -3,6 +3,7 @@ import itertools
import logging
import os
import platform
import shutil
import sys
import time
from dataclasses import dataclass
@ -128,14 +129,30 @@ class CliReport(Report):
py_exec_path: str = sys.executable
os_details: str = platform.platform()
_peak_memory_usage: int = 0
_peak_disk_usage: int = 0
def compute_stats(self) -> None:
mem_usage = psutil.Process(os.getpid()).memory_info().rss
if self._peak_memory_usage < mem_usage:
self._peak_memory_usage = mem_usage
self.peak_memory_usage = humanfriendly.format_size(self._peak_memory_usage)
try:
mem_usage = psutil.Process(os.getpid()).memory_info().rss
if self._peak_memory_usage < mem_usage:
self._peak_memory_usage = mem_usage
self.peak_memory_usage = humanfriendly.format_size(
self._peak_memory_usage
)
self.mem_info = humanfriendly.format_size(mem_usage)
disk_usage = shutil.disk_usage("/")
if self._peak_disk_usage < disk_usage.used:
self._peak_disk_usage = disk_usage.used
self.peak_disk_usage = humanfriendly.format_size(self._peak_disk_usage)
self.disk_info = {
"total": humanfriendly.format_size(disk_usage.total),
"used": humanfriendly.format_size(disk_usage.used),
"free": humanfriendly.format_size(disk_usage.free),
}
except Exception as e:
logger.warning(f"Failed to compute report memory usage: {e}")
self.mem_info = humanfriendly.format_size(mem_usage)
return super().compute_stats()