batch-shipyard/cascade/perf.py

121 строка
3.8 KiB
Python
Executable File

#!/usr/bin/env python3
# Copyright (c) Microsoft Corporation
#
# All rights reserved.
#
# MIT License
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.
# stdlib imports
import argparse
import datetime
import os
# non-stdlib imports
import azure.common
import azure.cosmosdb.table as azuretable
# global defines
_BATCHACCOUNT = os.environ['AZ_BATCH_ACCOUNT_NAME']
_POOLID = os.environ['AZ_BATCH_POOL_ID']
_NODEID = os.environ['AZ_BATCH_NODE_ID']
_PARTITION_KEY = '{}${}'.format(_BATCHACCOUNT, _POOLID)
def _create_credentials() -> azuretable.TableService:
"""Create storage credentials
:rtype: azure.cosmosdb.table.TableService
:return: azure storage table client
"""
sa, ep, sakey = os.environ['SHIPYARD_STORAGE_ENV'].split(':')
table_client = azuretable.TableService(
account_name=sa,
account_key=sakey,
endpoint_suffix=ep)
return table_client
def process_event(
table_client: azuretable.TableService,
table_name: str, source: str, event: str, ts: float,
message: str) -> None:
"""Process event
:param azure.cosmosdb.table.TableService table_client: table client
:param str table_name: table name
:param str source: source
:param str event: event
:param float ts: time stamp
:param str message: message
"""
entity = {
'PartitionKey': _PARTITION_KEY,
'RowKey': str(ts),
'Event': '{}:{}'.format(source, event),
'NodeId': _NODEID,
'Message': message,
}
while True:
try:
table_client.insert_entity(table_name, entity)
break
except azure.common.AzureConflictHttpError:
if not isinstance(ts, float):
ts = float(ts)
ts += 0.000001
entity['RowKey'] = str(ts)
def main():
"""Main function"""
# get command-line args
args = parseargs()
if args.ts is None:
args.ts = datetime.datetime.utcnow().timestamp()
args.source = args.source.lower()
args.event = args.event.lower()
# set up container name
table_name = args.prefix + 'perf'
# create storage credentials
table_client = _create_credentials()
# insert perf event into table
process_event(
table_client, table_name, args.source, args.event, args.ts,
args.message)
def parseargs():
"""Parse program arguments
:rtype: argparse.Namespace
:return: parsed arguments
"""
parser = argparse.ArgumentParser(
description='Batch Shipyard perf recorder')
parser.add_argument('source', help='event source')
parser.add_argument('event', help='event')
parser.add_argument('--ts', help='timestamp (posix)')
parser.add_argument('--message', help='message')
parser.add_argument('--prefix', help='storage container prefix')
return parser.parse_args()
if __name__ == '__main__':
main()