user/rubin/cpu.py

0001 #!/usr/bin/env python
0002 # Licensed under the Apache License, Version 2.0 (the "License");
0003 # you may not use this file except in compliance with the License.
0004 # You may obtain a copy of the License at
0005 # http://www.apache.org/licenses/LICENSE-2.0
0006 #
0007 # Authors:
0008 # - Paul Nilsson, paul.nilsson@cern.ch, 2021
0009
0010 from pilot.util.container import execute
0011
0012 import logging
0013 logger = logging.getLogger(__name__)
0014
0015
0016 def get_core_count(job):
0017     """
0018     Return the core count.
0019
0020     :param job: job object.
0021     :return: core count (int).
0022     """
0023
0024     return 0
0025
0026
0027 def add_core_count(corecount, core_counts=[]):
0028     """
0029     Add a core count measurement to the list of core counts.
0030
0031     :param corecount: current actual core count (int).
0032     :param core_counts: list of core counts (list).
0033     :return: updated list of core counts (list).
0034     """
0035
0036     return core_counts.append(corecount)
0037
0038
0039 def set_core_counts(job):
0040     """
0041     Set the number of used cores.
0042
0043     :param job: job object.
0044     :return:
0045     """
0046
0047     if job.pgrp:
0048         cmd = "ps axo pgid,psr | sort | grep %d | uniq | awk '{print $1}' | grep -x %d | wc -l" % (job.pgrp, job.pgrp)
0049         exit_code, stdout, stderr = execute(cmd, mute=True)
0050         logger.debug('%s: %s' % (cmd, stdout))
0051         try:
0052             job.actualcorecount = int(stdout)
0053         except Exception as e:
0054             logger.warning('failed to convert number of actual cores to int: %s' % e)
0055         else:
0056             logger.debug('set number of actual cores to: %d' % job.actualcorecount)
0057
0058             # overwrite the original core count and add it to the list
0059             job.corecount = job.actualcorecount
0060             job.corecounts = add_core_count(job.actualcorecount)
0061             logger.debug('current core counts list: %s' % str(job.corecounts))
0062
0063     else:
0064         logger.debug('payload process group not set - cannot check number of cores used by payload')