def test_supports_config():
expected_err_msg = "An error occurred while processing the parameter 'conflicting_types':\n" + \
"The specified type 'int' conflicts with the type of the default value " + \
r"'hello' \(type 'str'\)"
with pytest.raises(Exception, match=expected_err_msg):
@Source.supports_config(
Parameter('conflicting_types', type=int, default='hello'))
class _(Source.Source):
pass
import oauth2client
import pandas as pd
from packaging import version
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
from decisionengine_modules.util.retry_function import retry_wrapper
_MAX_RETRIES = 10
_RETRY_TIMEOUT = 10
@Source.supports_config(
Parameter("credential",
type=str,
comment="file that stores credentials in JSON"),
Parameter("max_retries", default=_MAX_RETRIES),
Parameter("retry_timeout", default=_RETRY_TIMEOUT),
)
@Source.produces(GCE_Occupancy=pd.DataFrame)
class GceOccupancy(Source.Source):
def __init__(self, config):
super().__init__(config)
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = config["credential"]
credentials, self.project = google.auth.default()
use_cache = version.parse(
oauth2client.__version__) < version.parse("4.0.0")
self.client = googleapiclient.discovery.build(
"compute",
from functools import partial
import logging
import pandas
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
from decisionengine_modules.util.retry_function import retry_wrapper
from decisionengine_modules.htcondor import htcondor_query
@Source.supports_config(
Parameter('condor_config',
type=str,
comment="path to condor configuration"),
Parameter('factories',
default=[],
comment="""Supported list entry layout:
{
'collector_host': 'factory_collector.com',
'classad_attrs': [],
'constraints': 'HTCondor classad query constraints'
}
"""), Parameter('nretries', default=0), Parameter('retry_interval', default=0))
@Source.produces(Factory_Entries_Grid=pandas.DataFrame,
Factory_Entries_AWS=pandas.DataFrame,
Factory_Entries_GCE=pandas.DataFrame,
Factory_Entries_LCF=pandas.DataFrame)
class FactoryEntries(Source.Source):
def __init__(self, config):
self.condor_config = config.get('condor_config')
import pandas as pd
from typing import Any
import decisionengine.framework.dataspace.datablock as datablock
import decisionengine.framework.dataspace.dataspace as dataspace
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
RETRIES = 10
RETRY_TO = 60
must_have = ('channel_name', 'Dataproducts')
@Source.supports_config(
Parameter('channel_name',
type=str,
comment="Channel from which to retrieve data products."),
Parameter('Dataproducts',
type=list,
comment="List of data products to retrieve."),
Parameter('retries',
default=RETRIES,
comment="Number of attempts allowed to fetch products."),
Parameter('retry_timeout',
default=RETRY_TO,
comment="Number of seconds to wait between retries."))
class SourceProxy(Source.Source):
def __init__(self, config):
if not set(must_have).issubset(set(config.keys())):
raise RuntimeError(
'SourceProxy misconfigured. Must have {} defined'.format(
# SPDX-FileCopyrightText: 2017 Fermi Research Alliance, LLC
# SPDX-License-Identifier: Apache-2.0
import pandas as pd
from bill_calculator_hep.GCEBillAnalysis import GCEBillCalculator
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
@Source.supports_config(
Parameter("projectId", type=str),
Parameter("credentialsProfileName", type=str, comment="not currently used"),
Parameter("accountNumber", type=int, comment="not currently used"),
Parameter("lastKnownBillDate", type=str, comment="in the form '%m/%d/%y %H:%M'"),
Parameter("balanceAtData", type=float, comment="in dollars"),
Parameter("applyDiscount", type=bool, comment="DLT discount does not apply to credits"),
Parameter("botoConfig", comment="not currently used"),
Parameter("localFileDir", type=dir, comment="location for downloaded billing files"),
)
@Source.produces(GCE_Billing_Info=pd.DataFrame)
class GCEBillingInfo(Source.Source):
def __init__(self, config):
super().__init__(config)
# Load configuration "constants"
self.projectId = config.get("projectId")
self.credentialsProfileName = config.get("credentialsProfileName") # NOT CURRENTLY USED
self.accountNumber = config.get("accountNumber") # NOT CURRENTLY USED
self.lastKnownBillDate = config.get("lastKnownBillDate") # '%m/%d/%y %H:%M'
self.balanceAtDate = config.get("balanceAtDate") # $
"""
This source takes input from instance_performance_gce.csv
and adds it to data block
"""
import pandas as pd
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
@Source.supports_config(
Parameter('csv_file', type=str, comment="path to CSV file"))
@Source.produces(GCE_Instance_Performance=pd.DataFrame)
class GCEInstancePerformance(Source.Source):
def __init__(self, config):
super().__init__(config)
self.csv_file = config.get('csv_file')
if not self.csv_file:
raise RuntimeError("No csv file found in configuration")
def acquire(self):
return {'GCE_Instance_Performance': pd.read_csv(self.csv_file)}
Source.describe(GCEInstancePerformance)
# SPDX-FileCopyrightText: 2017 Fermi Research Alliance, LLC
# SPDX-License-Identifier: Apache-2.0
import pandas as pd
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
@Source.supports_config(Parameter("multiplier", type=int))
@Source.supports_config(Parameter("channel_name", type=str))
@Source.produces(foo=pd.DataFrame)
class SourceWithSampleConfigNOP(Source.Source):
def __init__(self, config):
super().__init__(config)
self.multiplier = config.get("multiplier")
def acquire(self):
return {
"foo":
pd.DataFrame([
{
"col1": "value1",
"col2": 0.5 * self.multiplier
},
{
"col1": "value2",
"col2": 2.0 * self.multiplier
},
])
}
# SPDX-FileCopyrightText: 2017 Fermi Research Alliance, LLC
# SPDX-License-Identifier: Apache-2.0
import abc
import traceback
import pandas
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
from decisionengine_modules.htcondor import htcondor_query
@Source.supports_config(
Parameter("collector_host", type=str),
Parameter("condor_config", type=str),
Parameter("constraint", default=True),
Parameter("classad_attrs", type=list),
Parameter("group_attr", default=["Name"]),
Parameter("subsystem_name", type=str),
Parameter("correction_map", default={}),
)
class ResourceManifests(Source.Source, metaclass=abc.ABCMeta):
def __init__(self, config):
"""
In config files such as job_classification.jsonnet or Nersc.jsonnet,
put a dictionary named correction_map with keys corresponding to classad_attrs
and values that the operators want to be default values for the classad_attrs.
But here, we make available the option of an empty dictionary
because some classes that extend this class might not have correction_map
avaiable in its config file.
import pandas as pd
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
@Source.supports_config(Parameter('multiplier', type=int))
@Source.supports_config(Parameter('channel_name', type=str))
@Source.produces(foo=pd.DataFrame)
class SourceWithSampleConfigNOP(Source.Source):
def __init__(self, config):
super().__init__(config)
self.multiplier = config.get('multiplier')
def acquire(self):
return {
'foo':
pd.DataFrame([
{
'col1': 'value1',
'col2': 0.5 * self.multiplier
},
{
'col1': 'value2',
'col2': 2.0 * self.multiplier
},
])
}
Source.describe(SourceWithSampleConfigNOP,
import pandas as pd
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
@Source.supports_config(Parameter('multiplier', type=int))
@Source.produces(foo=pd.DataFrame)
class SourceWithSampleConfigNOP(Source.Source):
def __init__(self, config):
super().__init__(config)
self.multiplier = config.get('multiplier')
def acquire(self):
return {
'foo':
pd.DataFrame([
{
'col1': 'value1',
'col2': 0.5 * self.multiplier
},
{
'col1': 'value2',
'col2': 2.0 * self.multiplier
},
])
}
Source.describe(SourceWithSampleConfigNOP, sample_config={'multiplier': 1})
# SPDX-FileCopyrightText: 2017 Fermi Research Alliance, LLC
# SPDX-License-Identifier: Apache-2.0
import traceback
import pandas
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
from decisionengine_modules.htcondor import htcondor_query
@Source.supports_config(
Parameter("collector_host", type=str),
Parameter("schedds", default=[None]),
Parameter("condor_config", type=str),
Parameter("constraint", default=True),
Parameter("classad_attrs", type=list),
Parameter("correction_map", type=dict),
)
@Source.produces(job_manifests=pandas.DataFrame)
class JobQ(Source.Source):
def __init__(self, config):
"""
In config files such as job_classification.jsonnet or Nersc.jsonnet,
put a dictionary named correction_map with keys corresponding to classad_attrs
and values that the operators want to be default values for the classad_attrs.
"""
super().__init__(config)
self.collector_host = config.get("collector_host")
self.schedds = config.get("schedds", [None])
# SPDX-FileCopyrightText: 2017 Fermi Research Alliance, LLC
# SPDX-License-Identifier: Apache-2.0
"""
This dummy source takes the name of a source datablock
from config file as parameter "data_product_name" and produces
an empty pandas DataFrame as a datablock with that name
"""
import pandas as pd
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
@Source.supports_config(
Parameter("data_product_name", default=""), )
class EmptySource(Source.Source):
def __init__(self, config):
super().__init__(config)
self.data_product_name = config.get("data_product_name")
if not self.data_product_name:
raise RuntimeError("No data_product_name found in configuration")
self._produces = {self.data_product_name: pd.DataFrame}
def acquire(self):
self.logger.debug(f"in EmptySource: {self.data_product_name} acquire")
return {self.data_product_name: pd.DataFrame()}
Source.describe(EmptySource)
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
_SAMPLE_CONFIG = {
"target_aws_vm_burn_rate": [9.0],
"target_aws_bill_burn_rate": [10.0],
"target_aws_balance": [1000.0],
"target_gce_vm_burn_rate": [9.0],
"target_gce_balance": [1000.0]
}
@Source.supports_config(
Parameter('financial_parameters',
default=_SAMPLE_CONFIG,
comment='specifies the desired burn rates and balances'))
@Source.produces(financial_params=pandas.DataFrame)
class FinancialParameters(Source.Source):
def __init__(self, config):
super().__init__(config)
self.financial_parameters_dict = config.get('financial_parameters')
# The DataBlock given to the source is t=0
def acquire(self):
"""
Read the financial parameters from the config file and
return as a dataframe
"""
return {
'financial_parameters':
import logging
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
from decisionengine_modules.NERSC.util import newt
_MAX_RETRIES = 10
_RETRY_BACKOFF_FACTOR = 1
@Source.supports_config(Parameter('constraints',
type=dict,
comment="""Supports the layout:
{
'usernames': ['user1', 'user2'],
'newt_keys': {
'rname': ['m2612', 'm2696'],
'repo_type': ["STR", ],
}
"""),
Parameter('max_retries', default=_MAX_RETRIES),
Parameter('retry_backoff_factor', default=_RETRY_BACKOFF_FACTOR),
Parameter('passwd_file', type=str, comment="Path to password file"))
@Source.produces(Nersc_Allocation_Info=pd.DataFrame)
class NerscAllocationInfo(Source.Source):
"""
Information of allocations on NERSC machines
"""
# SPDX-FileCopyrightText: 2017 Fermi Research Alliance, LLC
# SPDX-License-Identifier: Apache-2.0
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
@Source.supports_config(Parameter("data_product_name", type=str))
class DynamicSource(Source.Source):
def __init__(self, config):
self.data_product_name = config["data_product_name"]
self._produces = {self.data_product_name: int}
def acquire(self):
return {self.data_product_name: 1}
Source.describe(DynamicSource)
# SPDX-FileCopyrightText: 2017 Fermi Research Alliance, LLC
# SPDX-License-Identifier: Apache-2.0
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
@Source.supports_config(Parameter("int_value", type=int))
@Source.produces(int_value=int)
class IntSource(Source.Source):
def __init__(self, config):
self._value = config["int_value"]
def acquire(self):
return {"int_value": self._value}
Source.describe(IntSource)
from functools import partial
import logging
import pandas
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
from decisionengine_modules.htcondor import htcondor_query
from decisionengine_modules.util.retry_function import retry_wrapper
@Source.supports_config(Parameter('condor_config', type=str),
Parameter('factories',
default=[],
comment="""Supported entries are of the form:
{
'collector_host': 'factory_collector-2.com',
'classad_attrs': [],
'constraints': 'HTCondor classad query constraints'
}
"""), Parameter('nretries', default=0), Parameter('retry_interval', default=0))
@Source.produces(factoryglobal_manifests=pandas.DataFrame)
class FactoryGlobalManifests(Source.Source):
def __init__(self, config):
if not config:
config = {}
if not isinstance(config, dict):
raise RuntimeError('parameters for module config should be a dict')
self.condor_config = config.get('condor_config')
import logging
import pandas as pd
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
from bill_calculator_hep.GCEBillAnalysis import GCEBillCalculator
@Source.supports_config(Parameter('projectId', type=str),
Parameter('credentialsProfileName', type=str, comment="not currently used"),
Parameter('accountNumber', type=int, comment="not currently used"),
Parameter('lastKnownBillDate', type=str, comment="in the form '%m/%d/%y %H:%M'"),
Parameter('balanceAtData', type=float, comment="in dollars"),
Parameter('applyDiscount', type=bool, comment="DLT discount does not apply to credits"),
Parameter('botoConfig', comment="not currently used"),
Parameter('localFileDir', type=dir, comment="location for downloaded billing files"))
@Source.produces(GCE_Billing_Info=pd.DataFrame)
class GCEBillingInfo(Source.Source):
def __init__(self, config):
super().__init__(config)
# Load configuration "constants"
self.projectId = config.get('projectId')
self.credentialsProfileName = config.get(
'credentialsProfileName') # NOT CURRENTLY USED
self.accountNumber = config.get('accountNumber') # NOT CURRENTLY USED
self.lastKnownBillDate = config.get(
'lastKnownBillDate') # '%m/%d/%y %H:%M'
self.balanceAtDate = config.get('balanceAtDate') # $
# Onix does not provide discounts
self.applyDiscount = config.get('applyDiscount')
from decisionengine.framework.modules.Source import Parameter
from decisionengine_modules.AWS.sources import DEAccountContants
from bill_calculator_hep.AWSBillAnalysis import AWSBillCalculator
@Source.supports_config(
Parameter(
'billing_configuration',
type=dict,
comment=
"""Configuration required to get AWS billing information. Supports the layout:
{
'AWSRnDAccountConstants': {
'lastKnownBillDate': '08/01/16 00:00', # '%m/%d/%y %H:%M'
'balanceAtDate': 3839.16, # $
'accountName': 'RnD',
'accountNumber': 159067897602,
'credentialsProfileName': 'BillingRnD',
'applyDiscount': True, # DLT discount does not apply to credits
'costRatePerHourInLastSixHoursAlarmThreshold': 2, # $ / h # $10/h
'costRatePerHourInLastDayAlarmThreshold': 2, # $ / h # $10/h
'emailReceipientForAlarms': '*****@*****.**'
}
}"""),
Parameter('dst_dir_for_s3_files',
type=str,
comment="Directory for AWS billing files"),
Parameter('verbose_flag', type=bool))
@Source.produces(AWS_Billing_Info=pd.DataFrame, AWS_Billing_Rate=pd.DataFrame)
class BillingInfo(Source.Source):
# SPDX-FileCopyrightText: 2017 Fermi Research Alliance, LLC
# SPDX-License-Identifier: Apache-2.0
from functools import partial
import pandas
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
from decisionengine_modules.htcondor import htcondor_query
from decisionengine_modules.util.retry_function import retry_wrapper
@Source.supports_config(
Parameter("condor_config", type=str),
Parameter(
"factories",
default=[],
comment="""Supported entries are of the form:
{
'collector_host': 'factory_collector-2.com',
'classad_attrs': [],
'constraints': 'HTCondor classad query constraints'
}
""",
),
Parameter("max_retries", default=0),
Parameter("retry_interval", default=0),
)
@Source.produces(factoryglobal_manifests=pandas.DataFrame)
for row in internal_data:
self.data.put(AvailabilityZone=row.AvailabilityZone,
InstanceType=row.InstanceType,
RunningVms=row.RunningVms)
'''
@Source.supports_config(
Parameter(
'occupancy_configuration',
type=str,
comment=
'''Python file containing (dynamic) account configuration. The format of the file is
{
"ProfileName1": ["RegionName1"],
"ProfileName2": ["RegionName2", RegionName3"]
...
}
where the keys ("ProfileName*") are the name of an account profile (e.g. "hepcloud-rnd"
and the entries in the lists (e.g. "RegionName1") are the name of a region (eg. "us-west-2").
'''))
@Source.produces(AWS_Occupancy=pd.DataFrame)
class AWSOccupancy(Source.Source):
def __init__(self, configdict):
self.config_file = configdict['occupancy_configuration']
self.logger = logging.getLogger()
def acquire(self):
"""
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
from decisionengine_modules.NERSC.util import newt
_MAX_RETRIES = 10
_RETRY_BACKOFF_FACTOR = 1
@Source.supports_config(
Parameter(
"constraints",
type=dict,
comment="""Supports the layout:
{
'usernames': ['user1', 'user2'],
'newt_keys': {
'rname': ['m2612', 'm2696'],
'repo_type': ["STR", ],
}
""",
),
Parameter("max_retries", default=_MAX_RETRIES),
Parameter("retry_backoff_factor", default=_RETRY_BACKOFF_FACTOR),
Parameter("passwd_file", type=str, comment="Path to password file"),
)
@Source.produces(Nersc_Allocation_Info=pd.DataFrame)
class NerscAllocationInfo(Source.Source):
"""
Information of allocations on NERSC machines
"""
import abc
import traceback
import pandas
import logging
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
from decisionengine_modules.htcondor import htcondor_query
@Source.supports_config(Parameter('collector_host', type=str),
Parameter('condor_config', type=str),
Parameter('constraint', default=True),
Parameter('classad_attrs', type=list),
Parameter('group_attr', default=['Name']),
Parameter('subsystem_name', type=str),
Parameter('correction_map', default={}))
class ResourceManifests(Source.Source, metaclass=abc.ABCMeta):
def __init__(self, config):
"""
In config files such as job_classification.jsonnet or Nersc.jsonnet,
put a dictionary named correction_map with keys corresponding to classad_attrs
and values that the operators want to be default values for the classad_attrs.
But here, we make available the option of an empty dictionary
because some classes that extend this class might not have correction_map
avaiable in its config file.
"""
super(Source.Source, self).__init__(config)
self.logger = logging.getLogger()
self.collector_host = config.get('collector_host')
self.condor_config = config.get('condor_config')
# SPDX-FileCopyrightText: 2017 Fermi Research Alliance, LLC
# SPDX-License-Identifier: Apache-2.0
"""
Fill in data from Instance Performance CSV file
"""
import pandas as pd
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
@Source.supports_config(Parameter("data_file", type=str, comment="CSV cost data file"))
@Source.produces(Performance_Data=pd.DataFrame)
class AWSInstancePerformance(Source.Source):
def __init__(self, config):
super().__init__(config)
self.data_file = config.get("data_file")
def acquire(self):
self.logger.debug("in AWSInstancePerformance acquire")
dataframe = (
pd.read_csv(self.data_file)
.drop_duplicates(subset=["AvailabilityZone", "InstanceType"], keep="last")
.reset_index(drop=True)
)
return {"Performance_Data": dataframe}
Source.describe(AWSInstancePerformance)
"""
Query Resource Limits from another channel with the factory source
"""
import typing
from decisionengine.framework.modules import Source, SourceProxy
from decisionengine.framework.modules.Source import Parameter
@Source.supports_config(Parameter('entry_limit_attrs', type=list))
@Source.produces(GCE_Resource_Limits=typing.Any)
class GCEResourceLimits(SourceProxy.SourceProxy):
"""
Consumes factory data to find GCE entry limits
"""
def __init__(self, config):
super().__init__(config)
self.entry_limit_attrs = config.get('entry_limit_attrs')
def acquire(self):
"""
Acquire google factory entry limits from source proxy
and return as pandas frame
:rtype: :obj:`~pd.DataFrame`
"""
factory_data = super().acquire()
df_factory_data = factory_data.get(self.data_keys[0])
df_entry_limits = df_factory_data[self.entry_limit_attrs]
return {'GCE_Resource_Limits': df_entry_limits}
"""
Fill in data from Job Limits CSV file
"""
import os
import time
import pandas as pd
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
_RETRIES = 5
_TO = 20
@Source.supports_config(
Parameter('data_file', type=str, comment="CSV job limits data file"))
@Source.produces(Job_Limits=pd.DataFrame)
class AWSJobLimits(Source.Source):
def __init__(self, config):
self.data_file = config['data_file']
def acquire(self):
rc = None
for i in range(_RETRIES):
if os.path.exists(self.data_file):
rc = {
'Job_Limits':
pd.read_csv(self.data_file).drop_duplicates(
subset=['AvailabilityZone', 'InstanceType'],
keep='last').reset_index(drop=True)
}
ll[i].data["Timestamp"] = spd.data["Timestamp"]
ll[i].data["SpotPrice"] = spd.data["SpotPrice"]
return ll
@Source.supports_config(
Parameter(
"spot_price_configuration",
type=str,
comment=
"""Python file containing (dynamic) account configuration. The format of the file is
config = {"ProfileName1":
{"RegionName1": ["Instance1", ], },
}
{
"ProfileName1": {"RegionName1": ["Instance1"]},
"ProfileName2": {"RegionName2": ["Instance2"]},
...
}
where the "ProfileName*" are names of account profiles (e.g. "hepcloud-rnd"),
the "RegionName* names of a regions (eg. "us-west-2"), and "Instance*" are
instance names. If the list of instances is empty, price information for
all instances is acquired.""",
))
@Source.produces(provisioner_resource_spot_prices=pd.DataFrame)
class AWSSpotPrice(Source.Source):
def __init__(self, config_dict):
super().__init__(config_dict)
self.config_file = config_dict["spot_price_configuration"]
import pandas as pd
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
from decisionengine_modules.NERSC.util import newt
_MAX_RETRIES = 10
_RETRY_BACKOFF_FACTOR = 1
# TODO this is a default column list and needs to be overriden from configuration
@Source.supports_config(Parameter('constraints', type=dict,
comment="""Supported dictionary structure:
{
machines: ["edison", "cori"],
newt_keys: {
user: ["user1", "user2"],
repo: ['m2612', 'm2696'],
features: ["knl&quad&cache"]
}
}"""),
Parameter('max_retries', default=_MAX_RETRIES),
Parameter('retry_backoff_factor', default=_RETRY_BACKOFF_FACTOR),
Parameter('passwd_file', type=str))
@Source.produces(Nersc_Job_Info=pd.DataFrame)
class NerscJobInfo(Source.Source):
"""
Information of jobs on NERSC machines
"""
def __init__(self, config):
super().__init__(config)
import logging
import pandas
import traceback
from decisionengine.framework.modules import Source
from decisionengine.framework.modules.Source import Parameter
from decisionengine_modules.htcondor import htcondor_query
@Source.supports_config(Parameter('collector_host', type=str),
Parameter('schedds', default=[None]),
Parameter('condor_config', type=str),
Parameter('constraint', default=True),
Parameter('classad_attrs', type=list),
Parameter('correction_map', type=dict))
@Source.produces(job_manifests=pandas.DataFrame)
class JobQ(Source.Source):
def __init__(self, config):
"""
In config files such as job_classification.jsonnet or Nersc.jsonnet,
put a dictionary named correction_map with keys corresponding to classad_attrs
and values that the operators want to be default values for the classad_attrs.
"""
super().__init__(config)
self.collector_host = config.get('collector_host')
self.schedds = config.get('schedds', [None])
self.condor_config = config.get('condor_config')
self.constraint = config.get('constraint', True)
self.classad_attrs = config.get('classad_attrs')
self.logger = logging.getLogger()
from decisionengine.framework.modules.Source import Parameter
from decisionengine_modules.NERSC.util import newt
_MAX_RETRIES = 10
_RETRY_BACKOFF_FACTOR = 1
# TODO this is a default column list and needs to be overriden from configuration
@Source.supports_config(
Parameter(
"constraints",
type=dict,
comment="""Supported dictionary structure:
{
machines: ["edison", "cori"],
newt_keys: {
user: ["user1", "user2"],
repo: ['m2612', 'm2696'],
features: ["knl&quad&cache"]
}
}""",
),
Parameter("max_retries", default=_MAX_RETRIES),
Parameter("retry_backoff_factor", default=_RETRY_BACKOFF_FACTOR),
Parameter("passwd_file", type=str),
)
@Source.produces(Nersc_Job_Info=pd.DataFrame)
class NerscJobInfo(Source.Source):
"""
Information of jobs on NERSC machines
"""