Jmuerto/yac deployment

src/user_workspaces_server/controllers/jobtypes/yac_job.py

@@ -0,0 +1,109 @@
+import logging
+import os
+from datetime import datetime, timedelta, timezone
+from urllib import parse
+
+import jwt
+import requests as http_r
+from django.apps import apps
+from django.template import loader
+
+from user_workspaces_server import models
+from user_workspaces_server.controllers.jobtypes.abstract_job import AbstractJob
+
+logger = logging.getLogger(__name__)
+
+
+class YACJob(AbstractJob):
+    def __init__(self, config, job_details):
+        super().__init__(config, job_details)
+        self.script_template_name = "yac_template.sh"
+
+    def get_script(self, template_params=None):
+        template_config = {"job_id": self.job_details["id"]}
+
+        # Generate JWT token for VITE authentication
+        if (jwt_secret_key := self.config.get("jwt_secret_key")) is None:
+            raise RuntimeError("jwt_secret_key is not set")
+
+        payload = {
+            "iat": datetime.now(timezone.utc),
+            "exp": datetime.now(timezone.utc) + timedelta(hours=6),
+        }
+        vite_auth_token = jwt.encode(payload, jwt_secret_key, algorithm="HS256")
+        template_config["vite_auth_token"] = vite_auth_token
+
+        template_config.update(self.config)
+        template_config.update(template_params)
+
+        template = loader.get_template(f"script_templates/{self.script_template_name}")
+        script = template.render(template_config)
+
+        return script
+
+    def status_check(self, job_model):
+        resource = apps.get_app_config("user_workspaces_server").main_resource
+
+        if job_model.status == models.Job.Status.FAILED:
+            return {
+                "message": "This job has failed. Support team has been notified and will investigate the error."
+            }
+
+        # Check to see if we already have a connection url in place.
+        if "connection_details" in job_model.job_details["current_job_details"]:
+            return {}
+
+        job_dir_path = os.path.join(
+            resource.resource_storage.root_dir,
+            job_model.workspace_id.file_path,
+            f".{job_model.id}",
+        )
+        # Open up the network_config file
+        subdomain = None
+        try:
+            with open(os.path.join(job_dir_path, ".network_config")) as f:
+                subdomain = f.readline().strip()
+                # TODO: Consider making the delimiter configurable
+                hostname, port = subdomain.split("-")
+                # We have to replace the periods with dashes for the dynamic naming
+                subdomain = subdomain.replace(".", "-")
+        except FileNotFoundError:
+            logger.warning("YAC network config missing.")
+            return {"current_job_details": {"message": "No network config found."}}
+
+        if not os.path.exists(os.path.join(job_dir_path, ".env")):
+            logger.warning(
+                f"Appyter output file {job_model.workspace_id.file_path}/.{job_model.id} missing."
+            )
+            return {"current_job_details": {"message": "Webserver not ready."}}
+
+        time_init = (
+            datetime.now(job_model.datetime_start.tzinfo) - job_model.datetime_start
+        ).total_seconds()
+
+        passthrough_url = parse.urlparse(resource.passthrough_domain)
+        url_domain = (
+            resource.passthrough_url
+            if subdomain is None
+            else f"{passthrough_url.scheme}://{subdomain}.{passthrough_url.netloc}"
+        )
+
+        # TODO: We need to turn off this verify False flag.
+        if http_r.get(url_domain, verify=False).status_code != 200:
+            logger.warning("Webserver not ready yet.")
+            return {"current_job_details": {"message": "Webserver not ready."}}
+
+        return {
+            "metrics": {
+                "time_init": time_init,
+            },
+            "current_job_details": {
+                "message": "Webserver ready.",
+                "proxy_details": {
+                    "hostname": hostname,
+                    "port": port,
+                    "path": "",
+                },
+                "connection_details": {"url_path": "", "url_domain": url_domain},
+            },
+        }

src/user_workspaces_server/templates/script_templates/yac_template.sh

@@ -0,0 +1,87 @@
+#!/bin/bash
+echo "STARTED @ $(date)"
+
+random_number () {
+    shuf -i ${1}-${2} -n 1
+}
+
+port_used_python() {
+  python -c "import socket; socket.socket().connect(('$1',$2))" >/dev/null 2>&1
+}
+
+port_used_python3() {
+  python3 -c "import socket; socket.socket().connect(('$1',$2))" >/dev/null 2>&1
+}
+
+port_used_nc(){
+  nc -w 2 "$1" "$2" < /dev/null > /dev/null 2>&1
+}
+
+port_used_lsof(){
+  lsof -i :"$2" >/dev/null 2>&1
+}
+
+port_used_bash(){
+  local bash_supported=$(strings /bin/bash 2>/dev/null | grep tcp)
+  if [ "$bash_supported" == "/dev/tcp/*/*" ]; then
+    (: < /dev/tcp/$1/$2) >/dev/null 2>&1
+  else
+    return 127
+  fi
+}
+
+# Check if port $1 is in use
+port_used () {
+  local port="${1#*:}"
+  local host=$((expr "${1}" : '\(.*\):' || echo "localhost") | awk 'END{print $NF}')
+  local port_strategies=(port_used_nc port_used_lsof port_used_bash port_used_python port_used_python3)
+
+  for strategy in ${port_strategies[@]};
+  do
+    $strategy $host $port
+    status=$?
+    if [[ "$status" == "0" ]] || [[ "$status" == "1" ]]; then
+      return $status
+    fi
+  done
+
+  return 127
+}
+
+# Find available port in range [$2..$3] for host $1
+# Default: [2000..65535]
+find_port () {
+  local host="${1:-localhost}"
+  local port=$(random_number "${2:-2000}" "${3:-65535}")
+  while port_used "${host}:${port}"; do
+    port=$(random_number "${2:-2000}" "${3:-65535}")
+  done
+  echo "${port}"
+}
+
+PORT=$(find_port)
+
+(
+umask 077
+cat > "$(pwd)/.network_config" << EOL
+$(hostname)-${PORT}
+EOL
+)
+
+(
+umask 077
+cat > "$(pwd)/.env" << EOL
+VITE_LLM_API_BASE_URL="{{ backend_url }}"
+VITE_DATA_PACKAGE_PATH="{{ data_manifest_path }}"
+VITE_PRODUCTION=true
+VITE_AUTH_TOKEN="{{ vite_auth_token }}"
+VITE_LLM_API_PORT=443
+EOL
+)
+
+mkdir "$(pwd)/build"
+
+# Launch the Apptainer YAC container
+set -x
+
+apptainer run --writable-tmpfs --bind "$(pwd)/build:/app/dist" --bind "$(pwd)/.env:/app/.env" --env YAC_PORT=${PORT} {{ sif_file_path }}

src/user_workspaces_server/utils.py

@@ -10,6 +10,7 @@ def translate_class_to_module(class_name):
         "JupyterLabJob": "jupyter_lab_job",
         "LocalTestJob": "local_test_job",
         "AppyterJob": "appyter_job",
+        "YACJob": "yac_job",
     }
 
     try: