diff --git a/cvat/apps/dataset_manager/task.py b/cvat/apps/dataset_manager/task.py index 7e1d4a14c2..cd98a082eb 100644 --- a/cvat/apps/dataset_manager/task.py +++ b/cvat/apps/dataset_manager/task.py @@ -22,7 +22,7 @@ from cvat.apps.engine.plugins import plugin_decorator from cvat.apps.engine.log import DatasetLogManager from cvat.apps.engine.utils import chunked_list from cvat.apps.events.handlers import handle_annotations_change -from cvat.apps.profiler import silk_profile +from cvat.apps.profiler import silk_profile, memray_tracker from cvat.apps.dataset_manager.annotation import AnnotationIR, AnnotationManager from cvat.apps.dataset_manager.bindings import TaskData, JobData, CvatImportError, CvatDatasetNotFoundError @@ -1103,6 +1103,7 @@ def export_task(task_id, dst_file, format_name, server_url=None, save_images=Fal with open(dst_file, 'wb') as f: task.export(f, exporter, host=server_url, save_images=save_images) +@memray_tracker @transaction.atomic def import_task_annotations(src_file, task_id, format_name, conv_mask_to_poly): task = TaskAnnotation(task_id) diff --git a/cvat/apps/profiler.py b/cvat/apps/profiler.py index 0a3885bb00..cc9239d747 100644 --- a/cvat/apps/profiler.py +++ b/cvat/apps/profiler.py @@ -1,3 +1,9 @@ +import hashlib +import subprocess +from datetime import datetime +from pathlib import Path + + from django.apps import apps if apps.is_installed("silk"): @@ -14,3 +20,92 @@ else: return wrapped return profile + + +class MemrayTracker: + def __init__(self, name: str) -> None: + self.name = name + self.tracker = None + self.trace_dir = None + + def __enter__(self): + import memray + from memray import FileFormat + + print("Initializing Memray Tracker") + current_revision = subprocess.check_output( + ["git", "rev-parse", "HEAD"], text=True + ).strip() + + traces_dir = ( + Path(__file__).parent.parent.parent + / "memray_traces" + / self.name + / current_revision[:8] +# / datetime.now().strftime("%Y%B%d-%I%M%p") + ) + + def get_hash_from_git_diff() -> str: + diff = subprocess.check_output(["git", "diff"], text=True).encode("utf-8") + return hashlib.sha256(diff).hexdigest() + + diff_hash = get_hash_from_git_diff() + self.trace_dir = traces_dir / f"{datetime.now().timestamp()}_{diff_hash}" + self.trace_dir.mkdir(parents=True, exist_ok=True) + + print(f"Will record traces to {self.trace_dir=}") + + _store_individual_file_diffs(self.trace_dir) + + self.tracker = memray.Tracker( + self.trace_dir / "trace.bin", + trace_python_allocators=True, + file_format=FileFormat.AGGREGATED_ALLOCATIONS, + ) + self.tracker.__enter__() + + def __exit__(self, exc_type, exc_val, exc_tb): + try: + return self.tracker.__exit__(exc_type, exc_val, exc_tb) + finally: + html_report = self.trace_dir / "report.html" + html_report_leaks = self.trace_dir / "report_leaks.html" + trace_file = self.trace_dir / "trace.bin" + subprocess.run( + ["memray", "flamegraph", "-o", str(html_report), str(trace_file)] + ) + subprocess.run(["open", str(html_report.absolute())]) + subprocess.run( + ["memray", "flamegraph", "--leaks", "-o", str(html_report_leaks), str(trace_file)] + ) + subprocess.run(["open", str(html_report_leaks.absolute())]) + + +def _store_individual_file_diffs(traces_dir: Path): + """Store the git diff for each individual file in the traces directory.""" + files_with_diffs = ( + subprocess.check_output(["git", "diff", "--name-only"], text=True) + .strip() + .splitlines() + ) + + for file_path in files_with_diffs: + diff_output = subprocess.check_output(["git", "diff", file_path], text=True) + file_name_safe = file_path.replace("/", "-") + patch_file_path = traces_dir / f"{file_name_safe}.patch" + patch_file_path.write_text(diff_output) + + +if apps.is_installed("memray"): + from functools import wraps + + def memray_tracker(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + with MemrayTracker(fn.__name__): + return fn(*args, **kwargs) + + return wrapper +else: + def memray_tracker(fn): + return fn diff --git a/cvat/settings/base.py b/cvat/settings/base.py index 404628fa55..12bbaef41c 100644 --- a/cvat/settings/base.py +++ b/cvat/settings/base.py @@ -117,6 +117,7 @@ INSTALLED_APPS = [ 'cvat.apps.events', 'cvat.apps.quality_control', 'cvat.apps.analytics_report', + 'memray', ] SITE_ID = 1