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))
def __init__(self, context: str):
self.context = context
self.config = client.Configuration()
self.loader = None
self.reload_task = None
self.api = None
self.core_v1 = None
self.client = None
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)
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
async def delete_deployment(self, deployment_id: str):
assert self.auth_client
assert self.cluster_endpoint
cfg = client.Configuration(
host=f"https://{self.cluster_endpoint}:443",
api_key={
"authorization": f"Bearer {await self.auth_client.get()}"
},
)
cfg.verify_ssl = False
async with ApiClient(configuration=cfg) as kube_api:
apps_api = client.AppsV1Api(kube_api)
core_api = client.CoreV1Api(kube_api)
# Delete service
service_id = f"{deployment_id}-svc"
await core_api.delete_namespaced_service(name=service_id,
namespace=KUBE_NAMESPACE)
# Delete deployment
await apps_api.delete_namespaced_deployment(
name=deployment_id, namespace=KUBE_NAMESPACE)
async def set_label(obj, hostnamePrefix, preemptible):
global NODEPOOL_LABEL
global CLUSTERS
body = {"metadata": {"labels": {NODEPOOL_LABEL: hostnamePrefix}}}
if preemptible:
# filter node if prepare-preemptible label not exist
if FILTER_PREEMPTIBLE_LABEL in obj['status']['nodeLabels']:
# preemptible exist
if PREEMPTIBLE_LABEL in obj['status']['nodeLabels']:
return
# add PREEMPTIBLE_LABEL if time WAIT_TIME_PREEMPTIBLE_LABEL
calculate_time = datetime.datetime.strptime(
obj['metadata']['creationTimestamp'],
"%Y-%m-%dT%H:%M:%SZ") + datetime.timedelta(
hours=WAIT_TIME_PREEMPTIBLE_LABEL)
if calculate_time < datetime.datetime.now():
body = {
"metadata": {
"labels": {
NODEPOOL_LABEL: hostnamePrefix,
PREEMPTIBLE_LABEL: "true"
}
},
"spec": {
"taints": [{
"effect":
"NoSchedule",
"key":
datetime.datetime.now().strftime('%s'),
"value":
"preemptible"
}]
}
}
print(
f"set {PREEMPTIBLE_LABEL}=true for node {obj['spec']['requestedHostname']}"
)
elif NODEPOOL_LABEL in obj['status']['nodeLabels']:
if obj['status']['nodeLabels'][NODEPOOL_LABEL] == hostnamePrefix:
return
print(
f"set {NODEPOOL_LABEL}={hostnamePrefix} for node {obj['spec']['requestedHostname']}"
)
# connect to cluster for this node
configuration = client.Configuration()
configuration.host = CLUSTERS[obj['metadata']['namespace']]['apiEndpoint']
configuration.ssl_ca_cert = FileOrData(
{
'certificate-authority':
CLUSTERS[obj['metadata']['namespace']]['caCert']
},
'certificate-authority',
data_key_name='certificate-authority').as_file()
configuration.api_key = {
"authorization":
"Bearer " +
CLUSTERS[obj['metadata']['namespace']]['serviceAccountToken']
}
client.Configuration.set_default(configuration)
v1 = client.CoreV1Api()
await v1.patch_node(obj['spec']['requestedHostname'],
body,
_request_timeout=30)
async def create_deployment(
self,
container: str,
num_replicas: int,
cpus: float = 1.0,
memory: float = 1.0,
) -> Tuple[str, str]:
assert self.auth_client
assert self.cluster_endpoint
cfg = client.Configuration(
host=f"https://{self.cluster_endpoint}:443",
api_key={
"authorization": f"Bearer {await self.auth_client.get()}"
},
)
cfg.verify_ssl = False
async with ApiClient(configuration=cfg) as kube_api:
apps_api = client.AppsV1Api(kube_api)
core_api = client.CoreV1Api(kube_api)
# Create deployment
deployment_id = f"dep-{uuid.uuid4()}"
deployment = client.V1Deployment(
api_version="apps/v1",
kind="Deployment",
metadata=client.V1ObjectMeta(name=deployment_id),
spec=client.V1DeploymentSpec(
replicas=num_replicas,
selector={"matchLabels": {
"dep": deployment_id
}},
template=client.V1PodTemplateSpec(
metadata=client.V1ObjectMeta(
labels={"dep": deployment_id}),
spec=client.V1PodSpec(containers=[
client.V1Container(
name=deployment_id,
env=[
client.V1EnvVar(name="PORT",
value=str(INTERNAL_PORT))
],
image=container,
resources=client.V1ResourceRequirements(
requests={
"cpu": str(cpus),
"memory": f"{int(memory * 1024)}M",
}),
ports=[
client.V1ContainerPort(
container_port=INTERNAL_PORT)
],
)
]),
),
),
)
await apps_api.create_namespaced_deployment(
namespace=KUBE_NAMESPACE, body=deployment)
# Create service
service_id = f"{deployment_id}-svc"
service_port = self.get_unassigned_port()
service = client.V1Service(
api_version="v1",
kind="Service",
metadata=client.V1ObjectMeta(
name=service_id,
# annotations={"cloud.google.com/load-balancer-type": "Internal"},
),
spec=client.V1ServiceSpec(
selector={"dep": deployment_id},
ports=[
client.V1ServicePort(
protocol="TCP",
port=service_port,
target_port=INTERNAL_PORT,
)
],
type="LoadBalancer",
),
)
await core_api.create_namespaced_service(namespace=KUBE_NAMESPACE,
body=service)
# Poll for external URL
service_ip = None
while not service_ip:
await asyncio.sleep(POLL_INTERVAL)
ingress = (await core_api.read_namespaced_service(
name=service_id,
namespace=KUBE_NAMESPACE)).status.load_balancer.ingress
if ingress:
service_ip = ingress[0].ip
service_url = f"http://{service_ip}:{service_port}"
print(f"Started deployment {deployment_id} at {service_url}")
return deployment_id, service_url