async def setup(self, app):
await super().setup(app)
try:
# Not a coroutine for some reason
config.load_incluster_config()
except config.ConfigException:
await config.load_kube_config()
self.api_client = client.ApiClient()
self.core_client = client.CoreV1Api(api_client=self.api_client)
self.custom_client = client.CustomObjectsApi(
api_client=self.api_client)
self.cluster_waiters = defaultdict(Flag)
self.clusters = {}
self.username_to_clusters = defaultdict(dict)
self.queue = WorkQueue()
self.informer = Informer(
parent=self,
name="cluster",
client=self.custom_client,
method="list_cluster_custom_object",
method_kwargs=dict(
group="gateway.dask.org",
version=self.crd_version,
plural="daskclusters",
label_selector=self.label_selector,
),
on_update=self.on_cluster_event,
on_delete=self.on_cluster_event,
)
await self.informer.start()
self.sync_task = asyncio.ensure_future(self.sync_clusters_loop())
async def simple_watch_clusters():
global CLUSTERS
this_clusters = {}
v1 = client.CustomObjectsApi()
async with watch.Watch().stream(v1.list_cluster_custom_object,
"management.cattle.io",
"v3",
"clusters",
timeout_seconds=10) as stream:
async for event in stream:
evt, obj = event['type'], event['object']
if obj['metadata']['name'] == 'local':
continue
cluster_id = obj['metadata']['name']
try:
credentials = {
'apiEndpoint': obj['status']['apiEndpoint'],
'caCert': obj['status']['caCert'],
'serviceAccountToken':
obj['status']['serviceAccountToken'],
}
this_clusters[cluster_id] = credentials
CLUSTERS[cluster_id] = credentials
except Exception as e:
print(f"Wait cluster {cluster_id}")
CLUSTERS = this_clusters
async def main():
# Setup Sentry if configured
sentry_dsn = os.getenv("SENTRY_DSN")
if sentry_dsn:
with open(".version") as f:
release = f.readline().strip()
environment = os.getenv("HOSTNAME", "dev").split("-")[0]
sentry_sdk.init(
sentry_dsn,
release=release,
environment=environment)
log.info("Sentry initialized with release='%s' and environment='%s'", release, environment)
try:
config.load_incluster_config()
except Exception:
await config.load_kube_config()
crds = client.CustomObjectsApi()
# Give tiller time to start up, if it isn't already
log.info("Waiting for tiller to be available ..")
await run_command(f"helm version", timeout=30)
tasks = [
asyncio.ensure_future(monitor_forever(crds, "global")),
asyncio.ensure_future(monitor_forever(crds, "production")),
asyncio.ensure_future(monitor_forever(crds, "staging")),
]
signal.signal(signal.SIGTERM, functools.partial(signal_handler, tasks))
await asyncio.wait(tasks)
async def api_client(context=None, api_client_kwargs=None):
await config.load_kube_config(config_file=KUBECONFIG_FILE)
context = context or {}
context['core_api'] = True
api_client_kwargs = api_client_kwargs or {}
api_client_kwargs.setdefault('request_timeout', 50)
api_cl = ApiClient(**api_client_kwargs)
user_context = {
'core_api': client.CoreV1Api(api_cl),
'apps_api': client.AppsV1Api(api_cl),
'storage_api': client.StorageV1Api(api_cl),
'batch_api': client.BatchV1Api(api_cl),
'cronjob_batch_api': client.BatchV1beta1Api(api_cl),
'custom_object_api': client.CustomObjectsApi(api_cl),
'extensions_api': client.ApiextensionsV1Api(api_cl),
}
try:
for k in filter(lambda k: context[k], context):
if k == 'node':
user_context[k] = await get_node(user_context['core_api'])
yield api_cl, user_context
finally:
await api_cl.close()
async def main():
try:
config.load_incluster_config()
except Exception:
await config.load_kube_config()
crds = client.CustomObjectsApi()
await monitor(crds)
def __init__(
self,
plural: str,
api_client: ApiClient,
queue: Queue[WatchEvent],
logger: BoundLogger,
) -> None:
self._plural = plural
self._queue = queue
self._logger = logger
self._api = client.CustomObjectsApi(api_client)
async def setup(self, k8s_config=None):
if not k8s_config:
if self.in_cluster:
config.load_incluster_config()
else:
await config.load_kube_config()
self.api_client = client.api_client.ApiClient()
else:
self.api_client = client.api_client.ApiClient(configuration=k8s_config)
self.k8s_api = client.CoreV1Api(self.api_client)
self.k8s_batch_api = client.BatchV1Api(self.api_client)
self.k8s_beta_api = client.ExtensionsV1beta1Api(self.api_client)
self.k8s_custom_object_api = client.CustomObjectsApi(self.api_client)
self.k8s_version_api = client.VersionApi(self.api_client)
async def get(cls) -> CustomObjectsApi:
if cls.k8s_custom_object_api:
return cls.k8s_custom_object_api
else:
try:
try:
await config.load_kube_config()
except FileNotFoundError:
config.load_incluster_config()
cls.k8s_custom_object_api = client.CustomObjectsApi(client.ApiClient())
return cls.k8s_custom_object_api
except Exception:
logger.exception(f'Failed to initialize {cls.__name__}')
raise
async def simple_watch_nodepools():
global NODEPOOLS
this_nodepools = {}
v1 = client.CustomObjectsApi()
async with watch.Watch().stream(v1.list_cluster_custom_object,
"management.cattle.io",
"v3",
"nodepools",
timeout_seconds=10) as stream:
async for event in stream:
evt, obj = event['type'], event['object']
nodepool_id = f"{obj['metadata']['namespace']}:{obj['metadata']['name']}"
hostnamePrefix = re.sub(r'([-_.])$', '',
obj['spec']['hostnamePrefix'])
this_nodepools[nodepool_id] = hostnamePrefix
NODEPOOLS[nodepool_id] = hostnamePrefix
NODEPOOLS = this_nodepools
async def simple_watch_nodes(preemptible=False):
global NODEPOOLS
v1 = client.CustomObjectsApi()
async with watch.Watch().stream(v1.list_cluster_custom_object,
"management.cattle.io",
"v3",
"nodes",
timeout_seconds=10) as stream:
async for event in stream:
evt, obj = event['type'], event['object']
if obj['spec']['nodePoolName'] and obj['spec'][
'nodePoolName'] in NODEPOOLS and evt in [
"ADDED", "MODIFIED"
]:
try:
await set_label(obj,
NODEPOOLS[obj['spec']['nodePoolName']],
preemptible)
except Exception as e:
# can't wait
print(f"Wait pool {obj['spec']['nodePoolName']}")
async def setup(self):
# Register signal handlers
loop = asyncio.get_event_loop()
for s in (signal.SIGTERM, signal.SIGINT):
loop.add_signal_handler(s, self.handle_shutdown_signal, s)
# Rate limiter for k8s api calls
self.rate_limiter = RateLimiter(rate=self.k8s_api_rate_limit,
burst=self.k8s_api_rate_limit_burst)
# Initialize the kubernetes clients
try:
config.load_incluster_config()
except config.ConfigException:
await config.load_kube_config()
self.api_client = client.ApiClient()
self.core_client = RateLimitedClient(
client.CoreV1Api(api_client=self.api_client), self.rate_limiter)
self.custom_client = RateLimitedClient(
client.CustomObjectsApi(api_client=self.api_client),
self.rate_limiter)
# Local state
self.cluster_info = collections.defaultdict(ClusterInfo)
self.stopped_clusters = {}
# Initialize queue and informers
self.queue = WorkQueue(
backoff=Backoff(base_delay=self.backoff_base_delay,
max_delay=self.backoff_max_delay))
endpoints_selector = (self.label_selector +
",app.kubernetes.io/component=dask-scheduler")
self.informers = {
"cluster":
Informer(
parent=self,
name="cluster",
client=self.custom_client,
method="list_cluster_custom_object",
method_kwargs=dict(
group="gateway.dask.org",
version=self.crd_version,
plural="daskclusters",
label_selector=self.label_selector,
),
on_update=self.on_cluster_update,
on_delete=self.on_cluster_delete,
),
"pod":
Informer(
parent=self,
name="pod",
client=self.core_client,
method="list_pod_for_all_namespaces",
method_kwargs=dict(label_selector=self.label_selector),
on_update=self.on_pod_update,
on_delete=self.on_pod_delete,
),
"endpoints":
Informer(
parent=self,
name="endpoints",
client=self.core_client,
method="list_endpoints_for_all_namespaces",
method_kwargs=dict(label_selector=endpoints_selector),
on_update=self.on_endpoints_update,
on_delete=self.on_endpoints_delete,
),
}
await asyncio.wait([i.start() for i in self.informers.values()])
self.log.debug("All informers started")
# Initialize reconcilers
self.reconcilers = [
asyncio.ensure_future(self.reconciler_loop())
for _ in range(self.parallelism)
]
# Start background tasks
self.task_pool = TaskPool()
self.task_pool.spawn(self.cleanup_expired_cluster_records_loop())
# Start the aiohttp application
self.runner = web.AppRunner(
self.app,
handle_signals=False,
access_log_class=AccessLogger,
access_log=self.log,
)
await self.runner.setup()
host, port = self.address.split(":")
port = int(port)
site = web.TCPSite(self.runner,
host,
port,
shutdown_timeout=15.0,
backlog=128)
await site.start()
self.log.info("%s started!", self.name)
self.log.info("API listening at http://%s", self.address)
def __init__(self, api_client: ApiClient, logger: BoundLogger) -> None:
self._api_client = api_client
self._api = client.CoreV1Api(api_client)
self._custom_api = client.CustomObjectsApi(api_client)
self._logger = logger