velorum/main.py at main · code0-tech/velorum · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
import concurrent.futures
import os
import signal
import sys
import time

import grpc

from src.interceptor.auth_interceptor import AuthInterceptor
from src.logger import get_logger

log = get_logger("velorum")

current_dir = os.path.dirname(os.path.abspath(__file__))
venv_site_packages = os.path.join(current_dir, ".venv", "lib", "python3.12", "site-packages")

# 2. Den Pfad hinzufügen, wo 'velorum' als Paket drinliegt
target_path = os.path.join(venv_site_packages, "tucana", "generated")

if target_path not in sys.path:
    sys.path.insert(0, target_path)

import tucana.generated.velorum.generate_pb2_grpc as generate_pb2_grpc
import tucana.generated.velorum.generate_pb2 as generate_pb2
import tucana.generated.velorum.info_pb2_grpc as info_pb2_grpc
import tucana.generated.velorum.info_pb2 as info_pb2
from grpc_reflection.v1alpha import reflection
from grpc_health.v1 import health, health_pb2, health_pb2_grpc

from src.endpoint.info.info_endpoint import InfoService
from src.endpoint.generation.generate_endpoint import GenerateService

if __name__ == '__main__':

    server = grpc.server(
        concurrent.futures.ThreadPoolExecutor(max_workers=10),
        interceptors=[AuthInterceptor()]
    )
    generate_pb2_grpc.add_GenerateServiceServicer_to_server(GenerateService(), server)
    info_pb2_grpc.add_InfoServiceServicer_to_server(InfoService(), server)


    health_servicer = health.HealthServicer()
    health_servicer.set('liveness', health_pb2.HealthCheckResponse.SERVING)
    health_servicer.set('readiness', health_pb2.HealthCheckResponse.SERVING)
    health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server)

    port = os.getenv("HOST", "0.0.0.0") + ":" + os.getenv("PORT", "50051")
    server.add_insecure_port(port)
    log.success(f"Velorum listening on {port}")  # type: ignore[attr-defined]

    SERVICE_NAMES = (
        generate_pb2.DESCRIPTOR.services_by_name['GenerateService'].full_name,
        info_pb2.DESCRIPTOR.services_by_name['InfoService'].full_name,
        reflection.SERVICE_NAME,
    )
    reflection.enable_server_reflection(SERVICE_NAMES, server)

    server.start()


    def shutdown(signum, frame):
        log.info("Shutting down gracefully...")
        health_servicer.enter_graceful_shutdown()
        server.stop(5).wait()

    signal.signal(signal.SIGTERM, shutdown)
    signal.signal(signal.SIGINT, shutdown)

    server.wait_for_termination()