Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
109 changes: 109 additions & 0 deletions src/user_workspaces_server/controllers/jobtypes/yac_job.py
Original file line number Diff line number Diff line change
@@ -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},
},
}
Original file line number Diff line number Diff line change
@@ -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 }}
1 change: 1 addition & 0 deletions src/user_workspaces_server/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
Loading