# Copyright 2018- The Pixie Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0

import px

#
# Get the Virtual memory usage and average memory for all processes in the k8s cluster.
#
t1 = px.DataFrame(table='process_stats', start_time='-30s')

t1.timestamp = px.bin(t1.time_, px.seconds(10))
t1.cmdline = px.upid_to_cmdline(t1.upid)

aggop = t1.groupby(['upid', 'timestamp', 'cmdline']).agg(
    vsize=('vsize_bytes', px.mean),
    rss=('rss_bytes', px.mean),
)

# Format column names.
aggop.pid = px.upid_to_pid(aggop.upid)
aggop.asid = px.upid_to_asid(aggop.upid)
aggop['Process Name'] = aggop.cmdline

# Uncomment and replace number to filter on a specific pid.
# aggop = aggop[aggop.pid == 4862]

aggop['Virtual Memory'] = aggop.vsize
aggop['Average Memory'] = aggop.rss
keep_columns = aggop[[
    'pid',
    'Process Name',
    'asid',
    'timestamp',
    'Virtual Memory',
    'Average Memory'
]]

px.display(keep_columns)
