async def watch_namespaces():
async with client.ApiClient() as api:
v1 = client.CoreV1Api(api)
async with watch.Watch().stream(v1.list_namespace) as stream:
async for event in stream:
etype, obj = event['type'], event['object']
print("{} namespace {}".format(etype, obj.metadata.name))
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 _watch(self):
DBSession = self.connector.DBSession
k8s_config.load_incluster_config()
async with k8s_client.ApiClient() as api:
v1 = k8s_client.CoreV1Api(api)
with open(os.path.join(self.config_dir, 'runner.namespace')) as fp:
namespace = fp.read().strip()
# Find existing run pods
pods = await v1.list_namespaced_pod(
namespace=namespace,
label_selector='app=run',
)
PROM_RUNS.set(0)
for pod in pods.items:
run_id = int(pod.metadata.labels['run'], 10)
logger.info("Found run pod for %d", run_id)
PROM_RUNS.inc()
await self._check_pod(api, run_id, pod)
# Watch changes
watch = k8s_watch.Watch()
f, kwargs = v1.list_namespaced_pod, dict(
namespace=namespace,
label_selector='app=run',
)
while True:
try:
async for event in watch.stream(f, **kwargs):
await self._handle_watch_event(api, DBSession, event)
except k8s_client.ApiException as e:
if e.status != 410:
raise
async def k8s() -> K8sAsync:
configuration = client.Configuration()
await config.load_kube_config(client_configuration=configuration)
api_client = client.ApiClient(configuration)
return K8sAsync(core_client=client.CoreApi(api_client),
v1_client=client.CoreV1Api(api_client))
async def initialize(self, logger: BoundLogger) -> None:
"""Initialize the dependency.
This must be called during application startup.
"""
await initialize_kubernetes()
self._api_client = client.ApiClient()
self._state = State()
async def watch_pods():
async with client.ApiClient() as api:
v1 = client.CoreV1Api(api)
async with watch.Watch().stream(
v1.list_pod_for_all_namespaces) as stream:
async for event in stream:
evt, obj = event['type'], event['object']
print("{} pod {} in NS {}".format(evt, obj.metadata.name,
obj.metadata.namespace))
def client(self) -> k8s_client.AppsV1Api:
if self._client is None:
a_configuration = k8s_client.Configuration(
host=self.configuration.host,
api_key={"authorization": self.configuration.token},
)
a_configuration.api_key_prefix["authorization"] = "Bearer"
a_configuration.verify_ssl = False
self._client = k8s_client.AppsV1Api(k8s_client.ApiClient(a_configuration))
return self._client
def main():
# Configs can be set in Configuration class directly or using helper
# utility. If no argument provided, the config will be loaded from
# default location.
await config.load_kube_config()
k8s_client = client.ApiClient()
await utils.create_from_yaml(k8s_client, "nginx-deployment.yaml")
k8s_api = client.ExtensionsV1beta1Api(k8s_client)
deps = await k8s_api.read_namespaced_deployment("nginx-deployment",
"default")
print("Deployment {0} created".format(deps.metadata.name))
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 init(cls, in_cluster: bool,
task_runner_service: TaskRunnerService) -> K8sClient:
if in_cluster:
# auth inside k8s cluster
config.load_incluster_config()
configuration = client.Configuration()
else:
# local auth (from kubectl config)
configuration = client.Configuration()
await config.load_kube_config(client_configuration=configuration)
api_client = client.ApiClient(configuration)
core_client = client.CoreApi(api_client)
v1_client = client.CoreV1Api(api_client)
return cls(core_client=core_client,
v1_client=v1_client,
task_runner_service=task_runner_service)
async def load(self) -> None:
self.loader = await config.load_kube_config(
context=self.context, client_configuration=self.config, persist_config=False
)
self.reload_task = asyncio.create_task(
config.refresh_token(self.loader, self.config)
)
self.api = client.ApiClient(configuration=self.config)
self.core_v1 = client.CoreV1Api(self.api)
# Build an HTTPX client that can talk to kube-apiserver.
client_cert = None
if self.config.cert_file:
client_cert = (self.config.cert_file, self.config.key_file)
ssl_context = httpx.create_ssl_context(
verify=self.config.verify_ssl,
cert=client_cert,
)
if self.config.ssl_ca_cert:
ssl_context.load_verify_locations(cafile=self.config.ssl_ca_cert)
self.client = httpx.AsyncClient(
auth=KubernetesAuth(self.config),
base_url=urljoin(self.config.host, "api/"),
verify=ssl_context,
)
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, return_type=None):
self._raw_return_type = return_type
self._stop = False
self._api_client = client.ApiClient()
self.resource_version = 0
self.resp = None
async def run_inner(self, run_info):
run_id = run_info['id']
del run_info
# This does not run the experiment, it schedules a runner pod by
# talking to the Kubernetes API. That pod will run the experiment and
# update the database directly
k8s_config.load_incluster_config()
name = self._pod_name(run_id)
# Load configuration from configmap volume
with open(os.path.join(self.config_dir, 'runner.pod_spec')) as fp:
pod_spec = yaml.safe_load(fp)
with open(os.path.join(self.config_dir, 'runner.namespace')) as fp:
namespace = fp.read().strip()
# Make required changes
for container in pod_spec['containers']:
if container['name'] == 'runner':
container['args'] += [str(run_id)]
# This is mostly used by Tilt
if os.environ.get('OVERRIDE_RUNNER_IMAGE'):
container['image'] = os.environ['OVERRIDE_RUNNER_IMAGE']
async with k8s_client.ApiClient() as api:
# Create a Kubernetes pod to run
v1 = k8s_client.CoreV1Api(api)
pod = k8s_client.V1Pod(
api_version='v1',
kind='Pod',
metadata=k8s_client.V1ObjectMeta(
name=name,
labels={
'app': 'run',
'run': str(run_id),
},
),
spec=pod_spec,
)
await v1.create_namespaced_pod(
namespace=namespace,
body=pod,
)
logger.info("Pod created: %s", name)
PROM_RUNS.inc()
# Create a service for proxy connections
svc = k8s_client.V1Service(
api_version='v1',
kind='Service',
metadata=k8s_client.V1ObjectMeta(
name=name,
labels={
'app': 'run',
'run': str(run_id),
},
),
spec=k8s_client.V1ServiceSpec(
selector={
'app': 'run',
'run': str(run_id),
},
ports=[
k8s_client.V1ServicePort(
protocol='TCP',
port=5597,
),
],
),
)
await v1.create_namespaced_service(
namespace=namespace,
body=svc,
)
logger.info("Service created: %s", name)
Convenience functions for creating pod templates.
"""
from collections import namedtuple
import copy
from kubernetes_asyncio import client
import json
try:
import yaml
except ImportError:
yaml = False
# FIXME: ApiClient provides us serialize / deserialize methods,
# but unfortunately also starts a threadpool for no reason! This
# takes up resources, so we try to not make too many.
SERIALIZATION_API_CLIENT = client.ApiClient()
def _set_k8s_attribute(obj, attribute, value):
"""
Set a specific value on a kubernetes object's attribute
obj
an object from Kubernetes Python API client
attribute
Should be a Kubernetes API style attribute (with camelCase)
value
Can be anything (string, list, dict, k8s objects) that can be
accepted by the k8s python client
"""
current_value = None
def __init__(self, return_type=None):
self._raw_return_type = return_type
self._stop = False
self._api_client = client.ApiClient()
