Skip to content

Instantly share code, notes, and snippets.

@duskvirkus
Last active November 25, 2021 20:40
Show Gist options
  • Save duskvirkus/648788e8005b109bd194958f7aa35179 to your computer and use it in GitHub Desktop.
Save duskvirkus/648788e8005b109bd194958f7aa35179 to your computer and use it in GitHub Desktop.
import time
import os
import threading
class ProfileResult:
def __init__(
self,
name: str,
start,
end,
process_id: int = 0,
thread_id: int = 0,
) -> None:
self.name: str = name
self.start = start
self.end = end
self.process_id = process_id
self.thread_id = thread_id
class Instrumentor:
def __init__(self) -> None:
self.file = None
self.profile_count: int = 0
def begin_session(
self,
name: str,
file_path: str = 'results.json'
) -> None:
self.file = open(file_path, "w")
self.write_header()
def end_session(self) -> None:
self.write_footer()
self.file.close()
self.session = None
self.profile_count = 0
def write_profile(
self,
result: ProfileResult,
) -> None:
if self.profile_count > 0:
self.file.write(',')
self.profile_count += 1
name: str = result.name
name = name.replace('"', "'")
self.file.write('{')
self.file.write('"cat":"function",')
self.file.write(f'"dur":{result.end - result.start},')
self.file.write(f'"name":"{name}",')
self.file.write('"ph":"X",')
self.file.write(f'"pid":{result.process_id},')
self.file.write(f'"tid":{result.thread_id},')
self.file.write(f'"ts":{result.start}')
self.file.write('}')
def write_header(self) -> None:
self.file.write('{"otherData": {},"traceEvents":[')
def write_footer(self) -> None:
self.file.write(']}')
class InstrumentationTimer:
def __init__(
self,
name: str,
process_id = None,
thread_id = None,
) -> None:
self.name = name
self.process_id = os.getpid()
if process_id is not None:
self.process_id = process_id
self.thread_id = threading.get_native_id()
if thread_id is not None:
self.thread_id = thread_id
self.stopped: bool = False
self.start = time.time_ns() / 1000
def __del__(self) -> None:
if not self.stopped:
self.stop()
def stop(self) -> None:
end = time.time_ns() / 1000
instance.write_profile(ProfileResult(self.name, self.start, end, self.process_id, self.thread_id))
self.stopped = True
instance = Instrumentor()
def configure_profiling(
save_location: str,
file_name: str = 'results',
) -> None:
instance.begin_session(os.path.join(save_location, file_name + '.json'))
def profile_decorator(func, pid = None, tid = None):
def wrapper(*args, **kwargs):
timer = InstrumentationTimer(func.__name__, pid, tid)
ret = func(*args, **kwargs)
timer.stop()
return ret
return wrapper
def end_profiling() -> None:
instance.end_session()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment