jupyter_notebook_locustfile.py

from locust import HttpUser, task, between, SequentialTaskSet, events
from locust.exception import StopUser
import time
import logging
from bs4 import BeautifulSoup
import csv
import itertools
import json
import random

# https://www.blazemeter.com/blog/how-to-run-locust-with-different-users
global USER_CREDENTIALS
USER_CREDENTIALS = None

with open("random_users.csv", "r") as f:
    reader = csv.reader(f)
    next(reader)
    USER_CREDENTIALS = list(reader)
    random.shuffle(USER_CREDENTIALS)


counter = itertools.count()
formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")


class DashboardTasks(SequentialTaskSet):
    lti_request_data = {}
    authenticity_token = None

    def on_start(self):
        self.logger = logging.getLogger(f"locust-{next(counter)}")
        self.logger.debug("Hatching locust")
        self.start_time = time.time()

        # NOTE here you set user specific values that will be used during authentication
        # NOTE this step will be specific to your authentication system
        if len(USER_CREDENTIALS) > 0:
            (
                self.user_number,
                self.gid,
                self.gidNumber,
                self.firstname,
                self.lastname,
            ) = USER_CREDENTIALS.pop(0)

            self.email = f"{self.firstname}{self.lastname}@gmail.com"

            self.lti_request_data = {
                "lti_message_type": "basic-lti-launch-request",
                "lti_version": "LTI-1p0",
                "oauth_consumer_key": "key",
                "resource_link_id": "resource",
                "lis_person_contact_email_primary": self.email,
                "lis_person_name_given": self.firstname,
                "lis_person_name_family": self.lastname,
                "custom_canvas_course_id": self.gid[2:],
                "custom_canvas_user_id": self.user_number,
            }

    # NOTE this is where you make the necessary requests to authenticate your user
    # NOTE this step will be specific to your authentication system
    @task
    def authenticate_user(self):

        # lti verification
        self.logger.debug("Posting to lti parse")
        response = self.client.post("/lti/launch", data=self.lti_request_data)

        # waiting page
        self.logger.debug("Visiting the waiting page")
        self.client.get(f"/verify_account_request/verify_account.html")

        # poll for success status of account before continuing, performing the same task that javascript would have on the waiting page
        while True:
            status_response = self.client.get(f"/verify_account_request/status")
            self.logger.info(status_response.json())

            if status_response.json()["status"] == "true":
                break

            time.sleep(10)

    @task
    def dashboard_page(self):
        self.logger.debug("Visiting the dashboard page")
        response = self.client.get(f"/pun/sys/dashboard")

        self.end_time = time.time()
        self.logger.info(
            f"Dashboard_Tasks_Elapsed_Time: {self.end_time - self.start_time}"
        )

        # while True:
        #     time.sleep(random.randint(50, 200))
        #     self.logger.debug("Visiting the dashboard page")
        #     response = self.client.get(f"/pun/sys/dashboard")

    @task
    def new_jupyter_page(self):
        self.logger.info("Visiting the jupyter notebook creation page")
        response = self.client.get(
            "/pun/sys/dashboard/batch_connect/sys/Jupyter/session_contexts/new"
        )

        soup = BeautifulSoup(response.content)

        # Get authenticity token necessary for creating notebook
        self.authenticity_token = soup.findAll("input", {"name": "authenticity_token"})[
            0
        ].get("value")
        self.logger.info(self.authenticity_token)

    @task
    def create_jupyter_notebook(self):
        form_data = {
            "utf8": "✓",
            "authenticity_token": self.authenticity_token,
            "batch_connect_session_context": {
                "custom_memory_per_node": "4",
                "custom_num_cores": "1",
                "custom_time": "04:00:00",
                "envscript": "",
                "bc_email_on_started": "0",
                "custom_email_address": "",
                "custom_reservation": "",
                "custom_nodename": "",
                "bc_account": "",
            },
            "commit": "Launch",
        }
        self.logger.info(form_data)

        # dispatch the jupyter notebook job
        response = self.client.post(
            "/pun/sys/dashboard/batch_connect/sys/Jupyter/session_contexts",
            json=(form_data),
        )

        # try to get the necessary values to login to the jupyter notebook
        # if the job is not yet running, the code will wait 10 seconds and try again
        while True:
            try:
                response = self.client.get("/pun/sys/dashboard/batch_connect/sessions")

                soup = BeautifulSoup(response.content)

                # Get password for accessing jupyter notebook
                self.jupyter_connect_form = soup.findAll(
                    "div", {"class": "ood-appkit markdown"}
                )[0].findChildren("form")[0]

                self.jupyter_connect_url = self.jupyter_connect_form.get("action")
                self.jupyter_dash_url = self.jupyter_connect_url.replace("login", "tree")
                self.logger.info(self.jupyter_connect_url)
                self.jupyter_connect_password = self.jupyter_connect_form.findChildren(
                    "input"
                )[0].get("value")

                self.logger.info(self.jupyter_connect_url)
                self.logger.info(self.jupyter_connect_password)

                break

            except Exception as e:
                time.sleep(10)
                pass

    @task
    def login_to_jupyter(self):

        response = self.client.post(
            self.jupyter_connect_url,
            data={"password": self.jupyter_connect_password},
            name="compute_node",
        )

        self.wait_end = time.time()
        self.logger.info(f"Time_To_Jupyter_Open: {self.wait_end - self.wait_start}")

    @task
    def refresh_jupyter(self):
        # simulate the user by refreshing the notebook every 30 to 90 seconds
        while True:
            time.sleep(random.randint(30, 90))
            self.logger.debug("Visiting jupyter notebook")
            response = self.client.get(self.jupyter_dash_url)

    @task
    def on_finish(self):
        raise StopUser()


class AcademicClusterUser(HttpUser):
    tasks = [DashboardTasks]
    wait_time = between(0, 0)
    host = "YOUR HOSTNAME"