initial commit

Dockerfile

@@ -0,0 +1,13 @@
+FROM python:3.9
+
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+WORKDIR /app
+
+COPY --chown=user ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY --chown=user . /app
+CMD ["uvicorn", "jupyter/jupyter_server:app", "--host", "0.0.0.0", "--port", "7860"]

jupyter/Dockerfile

@@ -0,0 +1,16 @@
+FROM python:3.9-slim
+
+WORKDIR /app
+
+# Minimal requirements
+RUN pip install --upgrade pip flask docker requests ipython jupyter-client ipykernel
+RUN ipython kernel install --name "python3" --user
+
+# Extra requirements
+RUN pip install pandas scikit-learn matplotlib seaborn
+
+COPY jupyter_kernel.py .
+
+EXPOSE 5000
+
+ENTRYPOINT ["python", "jupyter_kernel.py"]

jupyter/jupyter_kernel.py

@@ -0,0 +1,89 @@
+from flask import Flask, request, jsonify
+from jupyter_client import KernelManager
+import sys, io, time, queue, os
+
+
+app = Flask(__name__)
+TIMEOUT = os.getenv("TIMEOUT", 60)
+
+km = KernelManager(kernel_name='python3')
+
+km.start_kernel()
+kc = km.client()
+
+@app.route('/execute', methods=['POST'])
+def execute_code():
+    global kc
+    code = request.json.get('code')
+    if not code:
+        return jsonify({'error': 'No code provided'}), 400
+    
+    outputs = []
+    start_time = time.time()
+
+    _ = kc.execute(code)
+
+    success = True
+    error = None
+
+    while True:
+        try:
+            # if we timed out we interrupt the kernel so we can send the next request
+            if time.time() - start_time > TIMEOUT:
+                outputs.append({"text": f"Code execution has timed out (max {TIMEOUT}sec)."})
+                success = False
+                error = "TimeOut"
+                km.interrupt_kernel()
+                break
+            
+            # get message from output buffer
+            msg = kc.get_iopub_msg(timeout=1)
+
+            # if the kernel is idle again we can wrap up
+            if msg['header']['msg_type'] == 'status' and msg['content']['execution_state'] == 'idle':
+                break
+            
+            # save output messages
+            if msg['header']['msg_type'] == 'stream':
+                outputs.append({"text": msg['content']['text']})
+
+            # handle error messages
+            elif msg['header']['msg_type'] == 'error':
+                outputs.append({"text": "\n".join(msg['content']['traceback'])})
+                success = False
+                error = msg["content"]["ename"]
+
+            # handle figures
+            elif msg['header']['msg_type'] == 'display_data':
+                msg["content"]["data"]["text"] = msg["content"]["data"].pop("text/plain")
+                outputs.append(msg["content"]["data"])
+        
+        # if queue is empty we try again, unless the kernel is dead
+        except queue.Empty:
+            print("No message received (timeout)", )
+            if not kc.is_alive():
+                outputs.append({"text": "Kernel has died!"})
+                success = False
+                error = "KernelDied"
+                km.restart_kernel()
+                break
+
+    return jsonify({
+        'result': outputs,
+        'success': success,
+        'error': error
+    })
+
+@app.route('/restart', methods=['POST'])
+def restart_kernel():
+    global kc
+    km.restart_kernel()
+    kc = km.client()
+    return jsonify(status="Kernel has been restarted."), 200
+
+@app.route('/health', methods=['GET'])
+def health_check():
+    return jsonify(status="healthy"), 200
+
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=5000)

jupyter/jupyter_server.py

@@ -0,0 +1,159 @@
+from flask import Flask, request, jsonify
+import uuid
+import time
+import docker
+import requests
+import atexit
+import socket 
+import argparse
+import logging
+from pydantic import BaseModel, Field, ValidationError
+
+
+app = Flask(__name__)
+app.logger.setLevel(logging.INFO)
+
+
+# CLI function to parse arguments
+def parse_args():
+    parser = argparse.ArgumentParser(description="Jupyter server.")
+    parser.add_argument('--n_instances', type=int, help="Number of Jupyter instances.")
+    parser.add_argument('--n_cpus', type=int, default=2, help="Number of CPUs per Jupyter instance.")
+    parser.add_argument('--mem', type=str, default="2g", help="Amount of memory per Jupyter instance.")
+    parser.add_argument('--execution_timeout', type=int, default=10, help="Timeout period for a code execution.")
+    parser.add_argument('--port', type=int, default=5001, help="Port of main server")
+    return parser.parse_args()
+
+
+def get_unused_port(start=50000, end=65535, exclusion=[]):
+    for port in range(start, end + 1):
+        if port in exclusion:
+            continue
+        try:
+            sock = socket.socket()
+            sock.bind(("", port))
+            sock.listen(1)
+            sock.close()
+            return port
+        except OSError:
+            continue
+    raise IOError("No free ports available in range {}-{}".format(start, end))
+
+
+def create_kernel_containers(n_instances, n_cpus=2, mem="2g", execution_timeout=10):
+
+    docker_client = docker.from_env()
+    app.logger.info("Buidling docker image...")
+    image, logs = docker_client.images.build(path='./', tag='jupyter-kernel:latest')
+    app.logger.info("Building docker image complete.")
+
+    containers = []
+    port_exclusion = []
+    for i in range(n_instances):
+        
+        free_port =  get_unused_port(exclusion=port_exclusion)
+        port_exclusion.append(free_port) # it takes a while to startup so we don't use the same port twice
+        app.logger.info(f"Starting container {i} on port {free_port}...")
+        container = docker_client.containers.run(
+            "jupyter-kernel:latest",
+            detach=True,
+            mem_limit=mem,
+            cpuset_cpus=f"{i*n_cpus}-{(i+1)*n_cpus-1}",  # Limit to CPU cores 0 and 1
+            remove=True,
+            ports={'5000/tcp': free_port},
+            environment={"EXECUTION_TIMEOUT": execution_timeout},
+        )
+
+        containers.append({"container": container, "port": free_port})
+
+    start_time = time.time()
+    
+    containers_ready = []
+
+    while len(containers_ready) < n_instances:
+        app.logger.info("Pinging Jupyter containers to check readiness.")
+        if time.time() - start_time > 60:
+            raise TimeoutError("Container took too long to startup.")
+        for i in range(n_instances):
+            if i in containers_ready:
+                continue
+            url = f"http://localhost:{containers[i]['port']}/health"
+            try:
+                # TODO: dedicated health endpoint
+                response = requests.get(url)
+                if response.status_code == 200:
+                    containers_ready.append(i)
+            except Exception as e:
+                # Catch any other errors that might occur
+                pass
+        time.sleep(0.5)
+    app.logger.info("Containers ready!")
+    return containers
+
+def shutdown_cleanup():
+    app.logger.info("Shutting down. Stopping and removing all containers...")
+    for instance in app.containers:
+        try:
+            instance['container'].stop()
+            instance['container'].remove()
+        except Exception as e:
+            app.logger.info(f"Error stopping/removing container: {str(e)}")
+    app.logger.info("All containers stopped and removed.")
+
+
+class ServerRequest(BaseModel):
+    code: str = Field(..., example="print('Hello World!')")
+    instance_id: int = Field(0, example=0)
+    restart: bool = Field(False, example=False)
+
+
+@app.route('/execute', methods=['POST'])
+def execute_code():
+    try:
+        input = ServerRequest(**request.json)
+    except ValidationError as e:
+        return jsonify(e.errors()), 400
+
+
+    port = app.containers[input.instance_id]["port"]
+
+    app.logger.info(f"Received request for instance {input.instance_id} (port={port}).")
+
+    try:
+        if input.restart:
+            response = requests.post(f'http://localhost:{port}/restart', json={})
+            if response.status_code==200:
+                app.logger.info(f"Kernel for instance {input.instance_id} restarted.")
+            else:
+                app.logger.info(f"Error when restarting kernel of instance {input.instance_id}: {response.json()}.")
+
+        response = requests.post(f'http://localhost:{port}/execute', json={'code': input.code})
+        result = response.json()
+        return result
+    
+    except Exception as e:
+        app.logger.info(f"Error in execute_code: {str(e)}")
+        return jsonify({
+            'result': 'error',
+            'output': str(e)
+        }), 500
+
+
+atexit.register(shutdown_cleanup)
+
+if __name__ == '__main__':
+    args = parse_args()
+    app.containers = create_kernel_containers(
+        args.n_instances,
+        n_cpus=args.n_cpus, 
+        mem=args.mem, 
+        execution_timeout=args.execution_timeout
+        )
+    # don't use debug=True --> it will run main twice and thus start double the containers
+    app.run(debug=False, host='0.0.0.0', port=args.port) 
+
+
+# TODO:
+# how to mount data at runtime into the container? idea: mount a (read only) 
+# folder into the container at startup and copy the data in there. before starting 
+# the kernel we could cp the necessary data into the pwd.

requirements.txt

@@ -0,0 +1,5 @@
+flask
+jupyter-client
+pydantic
+docker
+uvicorn