mirror of
https://github.com/clearml/clearml-server
synced 2025-01-31 10:56:48 +00:00
127 lines
3.8 KiB
Python
127 lines
3.8 KiB
Python
from datetime import datetime
|
|
from os import getenv
|
|
|
|
from boltons.iterutils import first
|
|
from elasticsearch import Elasticsearch, Transport
|
|
|
|
from apiserver.config import config
|
|
|
|
log = config.logger(__file__)
|
|
|
|
OVERRIDE_HOST_ENV_KEY = (
|
|
"TRAINS_ELASTIC_SERVICE_HOST",
|
|
"ELASTIC_SERVICE_HOST",
|
|
"ELASTIC_SERVICE_SERVICE_HOST",
|
|
)
|
|
OVERRIDE_PORT_ENV_KEY = ("TRAINS_ELASTIC_SERVICE_PORT", "ELASTIC_SERVICE_PORT")
|
|
|
|
OVERRIDE_HOST = first(filter(None, map(getenv, OVERRIDE_HOST_ENV_KEY)))
|
|
if OVERRIDE_HOST:
|
|
log.info(f"Using override elastic host {OVERRIDE_HOST}")
|
|
|
|
OVERRIDE_PORT = first(filter(None, map(getenv, OVERRIDE_PORT_ENV_KEY)))
|
|
if OVERRIDE_PORT:
|
|
log.info(f"Using override elastic port {OVERRIDE_PORT}")
|
|
|
|
_instances = {}
|
|
|
|
|
|
class MissingClusterConfiguration(Exception):
|
|
"""
|
|
Exception when cluster configuration is not found in config files
|
|
"""
|
|
|
|
pass
|
|
|
|
|
|
class InvalidClusterConfiguration(Exception):
|
|
"""
|
|
Exception when cluster configuration does not contain required properties
|
|
"""
|
|
|
|
pass
|
|
|
|
|
|
class ESFactory:
|
|
@classmethod
|
|
def connect(cls, cluster_name):
|
|
"""
|
|
Returns the es client for the cluster.
|
|
Connects to the cluster if did not connect previously
|
|
:param cluster_name: Dot separated cluster path in the configuration file
|
|
:return: es client
|
|
:raises MissingClusterConfiguration: in case no config section is found for the cluster
|
|
:raises InvalidClusterConfiguration: in case cluster config section misses needed properties
|
|
"""
|
|
if cluster_name not in _instances:
|
|
cluster_config = cls.get_cluster_config(cluster_name)
|
|
hosts = cluster_config.get("hosts", None)
|
|
if not hosts:
|
|
raise InvalidClusterConfiguration(cluster_name)
|
|
|
|
args = cluster_config.get("args", {})
|
|
_instances[cluster_name] = Elasticsearch(
|
|
hosts=hosts, transport_class=Transport, **args
|
|
)
|
|
|
|
return _instances[cluster_name]
|
|
|
|
@classmethod
|
|
def get_all_cluster_names(cls):
|
|
return list(config.get("hosts.elastic"))
|
|
|
|
@classmethod
|
|
def get_cluster_config(cls, cluster_name):
|
|
"""
|
|
Returns cluster config for the specified cluster path
|
|
:param cluster_name: Dot separated cluster path in the configuration file
|
|
:return: config section for the cluster
|
|
:raises MissingClusterConfiguration: in case no config section is found for the cluster
|
|
"""
|
|
cluster_key = ".".join(("hosts.elastic", cluster_name))
|
|
cluster_config = config.get(cluster_key, None)
|
|
if not cluster_config:
|
|
raise MissingClusterConfiguration(cluster_name)
|
|
|
|
def set_host_prop(key, value):
|
|
for host in cluster_config.get("hosts", []):
|
|
host[key] = value
|
|
|
|
if OVERRIDE_HOST:
|
|
set_host_prop("host", OVERRIDE_HOST)
|
|
|
|
if OVERRIDE_PORT:
|
|
set_host_prop("port", OVERRIDE_PORT)
|
|
|
|
return cluster_config
|
|
|
|
@classmethod
|
|
def connect_all(cls):
|
|
clusters = config.get("hosts.elastic").as_plain_ordered_dict()
|
|
for name in clusters:
|
|
cls.connect(name)
|
|
|
|
@classmethod
|
|
def instances(cls):
|
|
return _instances
|
|
|
|
@classmethod
|
|
def timestamp_str_to_millis(cls, ts_str):
|
|
epoch = datetime.utcfromtimestamp(0)
|
|
current_date = datetime.strptime(ts_str, "%Y-%m-%dT%H:%M:%S.%fZ")
|
|
return int((current_date - epoch).total_seconds() * 1000.0)
|
|
|
|
@classmethod
|
|
def get_timestamp_millis(cls):
|
|
now = datetime.utcnow()
|
|
epoch = datetime.utcfromtimestamp(0)
|
|
return int((now - epoch).total_seconds() * 1000.0)
|
|
|
|
@classmethod
|
|
def get_es_timestamp_str(cls):
|
|
now = datetime.utcnow()
|
|
return now.strftime("%Y-%m-%dT%H:%M:%S") + ".%03d" % (now.microsecond / 1000) + "Z"
|
|
|
|
|
|
es_factory = ESFactory()
|