Enhance get_service_object & run_container_with_retries

.env.example

@@ -7,3 +7,11 @@ PREM_REGISTRY_URL=https://raw.githubusercontent.com/premAI-io/prem-registry/main
 # Sentry
 # ------------------------------------------------------------------------------------------
 SENTRY_DSN=https://[email protected]/4505244431941632
+
+# Proxy
+# ------------------------------------------------------------------------------------------
+PROXY_ENABLED=True
+
+# Dnsd
+# ------------------------------------------------------------------------------------------
+DNSD_URL=http://dnsd:8080

app/core/config.py

@@ -12,6 +12,11 @@
 SECRET_KEY: Secret = Secret(os.getenv("SECRET_KEY", ""))
 PROJECT_NAME: str = os.getenv("PROJECT_NAME", "Prem Daemon")
 
+# PROXY
+# ------------------------------------------------------------------------------
+PROXY_ENABLED: bool = os.getenv("PROXY_ENABLED", False)
+DNSD_URL: str = os.getenv("DNSD_URL", "http://dnsd:8080")
+
 # APIs
 # ------------------------------------------------------------------------------
 API_PREFIX = "/v1"
@@ -31,3 +36,11 @@
     level=logging.INFO,
     datefmt="%Y-%m-%d %H:%M:%S",
 )
+
+# Constants
+# ------------------------------------------------------------------------------
+DNSD_DNS_EXIST_PATH = "/dns/existing"
+
+
+def dns_exists_url() -> str:
+    return f"{DNSD_URL}{DNSD_DNS_EXIST_PATH}"

app/core/services.py

@@ -4,7 +4,7 @@
 import docker
 import psutil
 
-from app.core import utils
+from app.core import config, utils
 
 logger = logging.getLogger(__name__)
 
@@ -77,6 +77,13 @@ def get_service_object(
     else:
         service["downloaded"] = False
 
+    if config.PROXY_ENABLED:
+        domain = utils.check_dns_exists()
+
+        service["fullURL"] = f"{service['id']}.docker.localhost"
+        if domain:
+            service["fullURL"] = f"{service['id']}.{domain}"
+
     return service
 
 
@@ -149,9 +156,10 @@ def stop_all_running_services():
 
 def run_container_with_retries(service_object):
     client = utils.get_docker_client()
+    service_id = service_object["id"]  # Assuming the service ID is in service_object
 
     try:
-        client.containers.get(service_object["id"]).remove(force=True)
+        client.containers.get(service_id).remove(force=True)
     except Exception as error:
         logger.info(f"Failed to remove container {error}.")
 
@@ -167,7 +175,7 @@ def run_container_with_retries(service_object):
     volumes = {}
     if volume_path := service_object.get("volumePath"):
         try:
-            volume_name = f"prem-{service_object['id']}-data"
+            volume_name = f"prem-{service_id}-data"
             volume = client.volumes.create(name=volume_name)
             volumes = {volume.id: {"bind": volume_path, "mode": "rw"}}
         except Exception as error:
@@ -176,6 +184,22 @@ def run_container_with_retries(service_object):
     env_variables = service_object.get("envVariables", [])
     exec_commands = service_object.get("execCommands", [])
 
+    labels = {}
+    if config.PROXY_ENABLED:
+        dns_exists = utils.check_dns_exists()
+        if dns_exists:
+            labels = {
+                "traefik.enable": "true",
+                f"traefik.http.routers.{service_id}.rule": f"Host(`{service_id}.domain`)",
+                f"traefik.http.routers.{service_id}.entrypoints": "websecure",
+                f"traefik.http.routers.{service_id}.tls.certresolver": "myresolver",
+            }
+        else:
+            labels = {
+                "traefik.enable": "true",
+                f"traefik.http.routers.{service_id}.rule": f"Host({service_id}.docker.localhost)",
+            }
+
     for _ in range(10):
         try:
             container = client.containers.run(
@@ -184,10 +208,11 @@ def run_container_with_retries(service_object):
                 auto_remove=True,
                 detach=True,
                 ports={f"{service_object['defaultPort']}/tcp": port},
-                name=service_object["id"],
+                name=service_id,
                 volumes=volumes,
                 environment=env_variables,
                 device_requests=device_requests,
+                labels=labels,  # Add this line
             )
             logger.info(f"Started container {container.name}")
 

app/core/utils.py

@@ -196,3 +196,33 @@ def get_gpu_info():
     mem_percentage = (used_memory_value / total_memory_value) * 100
 
     return gpu_name, total_memory_value, used_memory_value, mem_percentage
+
+
+cached_domain = None
+
+
+def check_dns_exists():
+    global cached_domain
+
+    if cached_domain is not None:
+        return cached_domain
+
+    url = config.dns_exists_url()
+    try:
+        response = requests.get(url)
+        if response.status_code == 200 and response.content:
+            json_response = response.json()
+            if "domain" in json_response:
+                cached_domain = json_response["domain"]
+                return cached_domain
+            else:
+                print("Domain field not found in response.")
+                return None
+        else:
+            print(
+                f"Failed to get a valid response. Status Code: {response.status_code}"
+            )
+            return None
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        return None