micromegas 0.1.3__py3-none-any.whl → 0.1.7__py3-none-any.whl

micromegas/__init__.py

@@ -1,4 +1,10 @@
 import os
+from . import time
 from . import request
 from . import client
 from . import perfetto
+
+def connect():
+    "connect to the analytics service using default values"
+    BASE_URL = "http://localhost:8082/"
+    return client.Client(BASE_URL)

micromegas/client.py

@@ -1,4 +1,6 @@
 from . import request
+from . import time
+import cbor2
 
 
 class Client:
@@ -6,17 +8,28 @@ class Client:
         self.analytics_base_url = base_url + "analytics/"
         self.headers = headers
 
+    def find_process(self, process_id):
+        return request.request(
+            self.analytics_base_url + "find_process",
+            {"process_id": process_id},
+            headers=self.headers,
+        )
+
     def query_processes(self, begin, end, limit):
         return request.request(
             self.analytics_base_url + "query_processes",
-            {"begin": begin.isoformat(), "end": end.isoformat(), "limit": limit},
+            {
+                "begin": time.format_datetime(begin),
+                "end": time.format_datetime(end),
+                "limit": limit,
+            },
             headers=self.headers,
         )
 
     def query_streams(self, begin, end, limit, process_id=None, tag_filter=None):
         args = {
-            "begin": begin.isoformat(),
-            "end": end.isoformat(),
+            "begin": time.format_datetime(begin),
+            "end": time.format_datetime(end),
             "limit": limit,
             "process_id": process_id,
             "tag_filter": tag_filter,
@@ -30,8 +43,8 @@ class Client:
 
     def query_blocks(self, begin, end, limit, stream_id):
         args = {
-            "begin": begin.isoformat(),
-            "end": end.isoformat(),
+            "begin": time.format_datetime(begin),
+            "end": time.format_datetime(end),
             "limit": limit,
             "stream_id": stream_id,
         }
@@ -46,8 +59,8 @@ class Client:
         return request.request(
             self.analytics_base_url + "query_spans",
             {
-                "begin": begin.isoformat(),
-                "end": end.isoformat(),
+                "begin": time.format_datetime(begin),
+                "end": time.format_datetime(end),
                 "limit": limit,
                 "stream_id": stream_id,
             },
@@ -58,35 +71,100 @@ class Client:
         return request.request(
             self.analytics_base_url + "query_thread_events",
             {
-                "begin": begin.isoformat(),
-                "end": end.isoformat(),
+                "begin": time.format_datetime(begin),
+                "end": time.format_datetime(end),
                 "limit": limit,
                 "stream_id": stream_id,
             },
             headers=self.headers,
         )
-    
-    def query_log_entries(self, begin, end, limit, stream_id):
+
+    def query_log_entries(
+        self,
+        begin,
+        end,
+        limit=None,  # Necessary if stream_id is specified, ignored otherwise
+        stream_id=None,  # If none, query is run on cached lakehouse using query engine
+        sql=None,  # Necessary if stream_id is None, ignored otherwise
+    ):
         return request.request(
             self.analytics_base_url + "query_log_entries",
             {
-                "begin": begin.isoformat(),
-                "end": end.isoformat(),
+                "begin": time.format_datetime(begin),
+                "end": time.format_datetime(end),
                 "limit": limit,
                 "stream_id": stream_id,
+                "sql": sql,
             },
             headers=self.headers,
         )
 
-    def query_metrics(self, begin, end, limit, stream_id):
+    def query_metrics(self, begin, end, limit=None, stream_id=None, sql=None):
         return request.request(
             self.analytics_base_url + "query_metrics",
             {
-                "begin": begin.isoformat(),
-                "end": end.isoformat(),
+                "begin": time.format_datetime(begin),
+                "end": time.format_datetime(end),
                 "limit": limit,
                 "stream_id": stream_id,
+                "sql": sql,
             },
             headers=self.headers,
         )
-    
+
+    def query_view(self, view_set_name, view_instance_id, begin, end, sql):
+        return request.request(
+            self.analytics_base_url + "query_view",
+            {
+                "view_set_name": view_set_name,
+                "view_instance_id": view_instance_id,
+                "begin": time.format_datetime(begin),
+                "end": time.format_datetime(end),
+                "sql": sql,
+            },
+            headers=self.headers,
+        )
+
+    def query_partitions(self):
+        args = {}
+        return request.request(
+            self.analytics_base_url + "query_partitions",
+            args,
+            headers=self.headers,
+        )
+
+    def __stream_request(self, endpoint, args):
+        response = request.streamed_request(
+            self.analytics_base_url + endpoint,
+            args,
+            headers=self.headers,
+        )
+        while response.raw.readable():
+            try:
+                print(cbor2.load(response.raw))
+            except cbor2.CBORDecodeEOF:
+                break
+
+    def materialize_partitions(
+        self, view_set_name, view_instance_id, begin, end, partition_delta_seconds
+    ):
+        args = {
+            "view_set_name": view_set_name,
+            "view_instance_id": view_instance_id,
+            "begin": time.format_datetime(begin),
+            "end": time.format_datetime(end),
+            "partition_delta_seconds": partition_delta_seconds,
+        }
+        self.__stream_request("materialize_partitions", args)
+
+    def retire_partitions(
+        self, view_set_name, view_instance_id, begin, end, partition_delta_seconds
+    ):
+        args = {
+            "view_set_name": view_set_name,
+            "view_instance_id": view_instance_id,
+            "begin": time.format_datetime(begin),
+            "end": time.format_datetime(end),
+            "partition_delta_seconds": partition_delta_seconds,
+        }
+        self.__stream_request("retire_partitions", args)

micromegas/perfetto.py

@@ -1,21 +1,52 @@
 import crc
+from tqdm import tqdm
+
 
 # hack to allow perfetto proto imports
 # you can then import the protos like this: from protos.perfetto.trace import trace_pb2
 def load_perfetto_protos():
     import sys
     import pathlib
-    perfetto_folder =  pathlib.Path(__file__).parent.absolute() / "thirdparty/perfetto"
+
+    perfetto_folder = pathlib.Path(__file__).parent.absolute() / "thirdparty/perfetto"
     sys.path.append(str(perfetto_folder))
 
+
 def crc64_str(s):
     calculator = crc.Calculator(crc.Crc64.CRC64)
     return calculator.checksum(str.encode(s))
 
+
+def generate_batches(df_blocks):
+    nb_events_threshold = 1024 * 1024
+    begin = df_blocks.iloc[0]["begin_time"]
+    end = df_blocks.iloc[0]["end_time"]
+    nb_events = 0
+    for index, block in df_blocks.iterrows():
+        nb_events += block["nb_objects"]
+        end = block["end_time"]
+        if nb_events > nb_events_threshold:
+            yield (begin, end, nb_events)
+            begin = block["end_time"]
+            nb_events = 0
+    if nb_events > 0:
+        yield (begin, end, nb_events)
+
+
 class Writer:
-    def __init__( self, client, process_id, exe ):
+    """
+    Fetches thread events from the analytics server and formats them in the perfetto format.
+    Traces can be viewed using https://ui.perfetto.dev/
+    """
+
+    def __init__(self, client, process_id, exe):
         load_perfetto_protos()
         from protos.perfetto.trace import trace_pb2, trace_packet_pb2
+
+        self.names = {}
+        self.categories = {}
+        self.source_locations = {}
+        self.first = True
         self.client = client
         self.trace = trace_pb2.Trace()
         self.packets = self.trace.packet
@@ -27,8 +58,42 @@ class Writer:
         packet.track_descriptor.process.process_name = exe
         self.packets.append(packet)
 
-    def append_thread( self, begin, end, stream_id, thread_name, thread_id ):
+    def get_name_iid(self, name):
+        iid = self.names.get(name)
+        is_new = False
+        if iid is None:
+            is_new = True
+            iid = len(self.names) + 1
+            self.names[name] = iid
+        return iid, is_new
+
+    def get_category_iid(self, cat):
+        iid = self.categories.get(cat)
+        is_new = False
+        if iid is None:
+            is_new = True
+            iid = len(self.categories) + 1
+            self.categories[cat] = iid
+        return iid, is_new
+
+    def get_location_iid(self, loc):
+        iid = self.source_locations.get(loc)
+        is_new = False
+        if iid is None:
+            is_new = True
+            iid = len(self.source_locations) + 1
+            self.source_locations[loc] = iid
+        return iid, is_new
+    
+    def append_thread(self, stream_id, thread_name, thread_id):
         from protos.perfetto.trace import trace_pb2, trace_packet_pb2, track_event
+
+        df_blocks = self.client.query_blocks(
+            begin=None, end=None, limit=100_000, stream_id=stream_id
+        )
+        if df_blocks.empty:
+            return
+
         packet = trace_packet_pb2.TracePacket()
         thread_uuid = crc64_str(stream_id)
         packet.track_descriptor.uuid = thread_uuid
@@ -39,26 +104,102 @@ class Writer:
         self.packets.append(packet)
         trusted_packet_sequence_id = 1
 
-        df_events = self.client.query_thread_events(begin, end, limit=1024*1024, stream_id = stream_id)
-        df_events["ns"] = df_events["timestamp"].astype('int64')
-        for index, event in df_events.iterrows():
-            packet = trace_packet_pb2.TracePacket()
-            packet.timestamp = event["ns"]
-            if event["event_type"] == "begin":
-                packet.track_event.type = track_event.track_event_pb2.TrackEvent.Type.TYPE_SLICE_BEGIN
-            elif event["event_type"] == "end":
-                packet.track_event.type = track_event.track_event_pb2.TrackEvent.Type.TYPE_SLICE_END
-            else:
-                raise Exception("unknown event type")
-            packet.track_event.track_uuid = thread_uuid
-            packet.track_event.name = event["name"]
-            packet.trusted_packet_sequence_id = trusted_packet_sequence_id
-            self.packets.append(packet)
-
-    def write_file( self, filename ):
+        batches = list(generate_batches(df_blocks))
+        for begin, end, limit in tqdm(batches, unit="event batches"):
+            df_spans = self.client.query_spans(
+                begin, end, limit=limit, stream_id=stream_id
+            )
+            begin_ns = df_spans["begin"].astype("int64")
+            end_ns = df_spans["end"].astype("int64")
+            for index, span in df_spans.iterrows():
+                packet = trace_packet_pb2.TracePacket()
+                packet.timestamp = begin_ns[index]
+                packet.track_event.type = (
+                    track_event.track_event_pb2.TrackEvent.Type.TYPE_SLICE_BEGIN
+                )
+                packet.track_event.track_uuid = thread_uuid
+                span_name = span["name"]
+                name_iid, new_name = self.get_name_iid(span_name)
+                packet.track_event.name_iid = name_iid
+                category_iid, new_category = self.get_category_iid(span["target"])
+                packet.track_event.category_iids.append(category_iid)
+
+                source_location = (span["filename"], span["line"])
+                source_location_iid, new_source_location = self.get_location_iid(source_location)
+                packet.track_event.source_location_iid = source_location_iid
+                if self.first:
+                    # this is necessary for interning to work
+                    self.first = False
+                    packet.first_packet_on_sequence = True
+                    packet.sequence_flags = 3
+                else:
+                    packet.sequence_flags = 2
+
+                if new_name:
+                    event_name = packet.interned_data.event_names.add()
+                    event_name.iid = name_iid
+                    event_name.name = span_name
+                if new_category:
+                    cat_name = packet.interned_data.event_categories.add()
+                    cat_name.iid = category_iid
+                    cat_name.name = span["target"]
+                if new_source_location:
+                    loc = packet.interned_data.source_locations.add()
+                    loc.iid = source_location_iid
+                    loc.file_name = source_location[0]
+                    loc.line_number = source_location[1]
+
+                packet.trusted_packet_sequence_id = trusted_packet_sequence_id
+                self.packets.append(packet)
+
+                packet = trace_packet_pb2.TracePacket()
+                packet.timestamp = end_ns[index]
+                packet.track_event.type = (
+                    track_event.track_event_pb2.TrackEvent.Type.TYPE_SLICE_END
+                )
+                packet.track_event.track_uuid = thread_uuid
+                packet.track_event.name_iid = name_iid
+                packet.track_event.category_iids.append(category_iid)
+                packet.track_event.source_location_iid = source_location_iid
+                packet.sequence_flags = 2
+                packet.trusted_packet_sequence_id = trusted_packet_sequence_id
+
+                self.packets.append(packet)
+
+    def write_file(self, filename):
         with open(filename, "wb") as f:
             f.write(self.trace.SerializeToString())
-        
-        
 
-        
+
+def get_process_cpu_streams(client, process_id):
+    def prop_to_dict(props):
+        prop_dict = {}
+        for p in props:
+            prop_dict[p["key"]] = p["value"]
+        return prop_dict
+
+    def get_thread_name(prop_dict):
+        return prop_dict["thread-name"]
+
+    def get_thread_id(prop_dict):
+        return int(prop_dict["thread-id"])
+
+    df_streams = client.query_streams(
+        begin=None, end=None, limit=1024, tag_filter="cpu", process_id=process_id
+    )
+    df_streams["properties"] = df_streams["properties"].apply(prop_to_dict)
+    df_streams["thread_name"] = df_streams["properties"].apply(get_thread_name)
+    df_streams["thread_id"] = df_streams["properties"].apply(get_thread_id)
+    return df_streams
+
+
+def write_process_trace(client, process_id, trace_filepath):
+    process_df = client.find_process(process_id)
+    assert process_df.shape[0] == 1
+    process = process_df.iloc[0]
+    streams = get_process_cpu_streams(client, process_id)
+    writer = Writer(client, process_id, process["exe"])
+    for index, stream in tqdm(list(streams.iterrows()), unit="threads"):
+        stream_id = stream["thread_id"]
+        writer.append_thread(stream["stream_id"], stream["thread_name"], stream_id)
+    writer.write_file(trace_filepath)

micromegas/request.py

@@ -18,3 +18,19 @@ def request(url, args, headers={}):
         )
     table = pq.read_table(io.BytesIO(response.content))
     return table.to_pandas()
+
+def streamed_request(url, args, headers={}):
+    response = requests.post(
+        url,
+        headers=headers,
+        data=cbor2.dumps(args),
+        stream=True,
+        timeout=300,
+    )
+    if response.status_code != 200:
+        raise Exception(
+            "http request url={2} failed with code={0} text={1}".format(
+                response.status_code, response.text, url
+            )
+        )
+    return response

micromegas/time.py

@@ -0,0 +1,30 @@
+import datetime
+import pandas
+import re
+
+def format_datetime(value):
+    nonetype = type(None)
+    match type(value):
+        case datetime.datetime:
+            if value.tzinfo is None:
+                raise RuntimeError("datetime needs a valid time zone")
+            return value.isoformat()
+        case pandas.Timestamp:
+            return value.isoformat()
+        case nonetype:
+            return None
+    raise RuntimeError("value of unknown type in format_datetime")
+
+def parse_time_delta(user_string):
+    parser = re.compile("(\\d+)([mhd])")
+    m = parser.match(user_string)
+    nbr = int(m.group(1))
+    unit = m.group(2)
+    if unit == "m":
+        return datetime.timedelta(minutes=nbr)
+    elif unit == "h":
+        return datetime.timedelta(hours=nbr)
+    elif unit == "d":
+        return datetime.timedelta(days=nbr)
+    else:
+        raise RuntimeError("invalid time delta: " + user_string)

{micromegas-0.1.3.dist-info → micromegas-0.1.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: micromegas
-Version: 0.1.3
+Version: 0.1.7
 Summary: Python analytics client for https://github.com/madesroches/micromegas/
 Author: Marc-Antoine Desroches
 Author-email: madesroches@gmail.com
@@ -16,6 +16,7 @@ Requires-Dist: protobuf (>=5.27.1,<6.0.0)
 Requires-Dist: pyarrow (>=16.0.0,<17.0.0)
 Requires-Dist: requests (>=2.31.0,<3.0.0)
 Requires-Dist: tabulate (>=0.9.0,<0.10.0)
+Requires-Dist: tqdm (>=4.66.5,<5.0.0)
 Description-Content-Type: text/markdown
 
 # Micromegas

{micromegas-0.1.3.dist-info → micromegas-0.1.7.dist-info}/RECORD

@@ -1,7 +1,7 @@
-micromegas/__init__.py,sha256=zIsOSogb45OohIYwtNrXF0tocz2PvbvTolf979tSrtU,76
-micromegas/client.py,sha256=TlwlAgk5E6A6E4DfwSClg1H5trbwhupfscgHqqWoqmA,2783
-micromegas/perfetto.py,sha256=F7cB-iQjk35EBooIXmtgSVntjt00jgubFwnriJGl5_I,2697
-micromegas/request.py,sha256=u6nqXV5iQjLdXjVrzDqrDV1yoqUbeJ9qjiVgz87-uN8,512
+micromegas/__init__.py,sha256=E_j3LFxMk9rSMJunwDCi_90NsRHm1fKwjj_6KGMYCjQ,246
+micromegas/client.py,sha256=2ejIRukZiJp6Q8YncK9dJtaazXX1s5TKNaKzndRxswk,5509
+micromegas/perfetto.py,sha256=yuIe5iKvca61aWMBQNziSGM-DHcOEsiobtKx2SsNQ3E,7829
+micromegas/request.py,sha256=NV0urom5P3_P2q94gX51hxW_Fnrp_DDRorsP3mUb5NM,941
 micromegas/thirdparty/perfetto/protos/perfetto/common/android_energy_consumer_descriptor_pb2.py,sha256=l8QNXqnB-mJIkuFr2s1YoLQXHm3G-ZcOGp_OW_hQ0TE,1887
 micromegas/thirdparty/perfetto/protos/perfetto/common/android_log_constants_pb2.py,sha256=O5zDZkV8Nji0O2ryJRP4FTWdgdOBlDymWNcpNNDOFxk,2017
 micromegas/thirdparty/perfetto/protos/perfetto/common/builtin_clock_pb2.py,sha256=7qLL_BENTxRFQH8DfHDvyWAkgwy0VHrOaE8XhL8iZgk,1822
@@ -207,6 +207,7 @@ micromegas/thirdparty/perfetto/protos/perfetto/trace/track_event/track_event_pb2
 micromegas/thirdparty/perfetto/protos/perfetto/trace/translation/translation_table_pb2.py,sha256=-hkUdv07TsSDHH1mier2KyAhmivK4GSzEfAzAEYv20U,6630
 micromegas/thirdparty/perfetto/protos/perfetto/trace/trigger_pb2.py,sha256=We7Yi8o3cEcrSNxY1zLUUO6tEWnD36C2f3O_s8_qv0I,1435
 micromegas/thirdparty/perfetto/protos/perfetto/trace/ui_state_pb2.py,sha256=Af-SXwhroNhRXMrtw6e2eU1liCImMRxSdmkt_AuSHf8,1752
-micromegas-0.1.3.dist-info/METADATA,sha256=g_S3KRAQOMELgMcL8lfFoMXuBTvPSEt3qN_L_SexKLY,839
-micromegas-0.1.3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-micromegas-0.1.3.dist-info/RECORD,,
+micromegas/time.py,sha256=teMWk_hniW2jI7MWJ2w0HuckdTiebdr35-snUz_3cfU,911
+micromegas-0.1.7.dist-info/METADATA,sha256=c7Dyj3APPSt59QLCF8id6Hcp6rjUDT_owecyrweWhAM,877
+micromegas-0.1.7.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+micromegas-0.1.7.dist-info/RECORD,,