diff --git a/.github/workflows/ci_frontend.yml b/.github/workflows/ci_frontend.yml index 92bfd047..554aaea9 100644 --- a/.github/workflows/ci_frontend.yml +++ b/.github/workflows/ci_frontend.yml @@ -17,4 +17,4 @@ jobs: - name: Launch frontend tests run: | . ./env.sh - docker-compose run clockwork_dev scripts/launch_frontend_tests_in_clockwork_dev.sh + docker compose run clockwork_dev scripts/launch_frontend_tests_in_clockwork_dev.sh diff --git a/clockwork_frontend_test/test_jobs_search.py b/clockwork_frontend_test/test_jobs_search.py index db5cd55e..a89fb7fd 100644 --- a/clockwork_frontend_test/test_jobs_search.py +++ b/clockwork_frontend_test/test_jobs_search.py @@ -642,7 +642,7 @@ def test_filter_by_job_array(page: Page): def test_filter_by_job_user_props(page: Page): # Login - page.goto(f"{BASE_URL}/login/testing?user_id=student00@mila.quebec") + page.goto(f"{BASE_URL}/login/testing?user_id=student01@mila.quebec") # Go to settings. page.goto(f"{BASE_URL}/settings/") radio_job_user_props = page.locator("input#jobs_list_job_user_props_toggle") diff --git a/clockwork_frontend_test/test_settings.py b/clockwork_frontend_test/test_settings.py index 39d58045..356bdc37 100644 --- a/clockwork_frontend_test/test_settings.py +++ b/clockwork_frontend_test/test_settings.py @@ -227,7 +227,19 @@ def test_jobs_search_columns(page: Page): def test_jobs_search_column_job_user_props(page: Page): # Login - page.goto(f"{BASE_URL}/login/testing?user_id=student00@mila.quebec") + page.goto(f"{BASE_URL}/login/testing?user_id=student01@mila.quebec") + + # Change language in settings + page.goto(f"{BASE_URL}/settings/") + # Get language select. + select = page.locator("select#language_selection") + # Check default language is French. + expect(select).to_have_value("fr") + # Switch to english. + select.select_option("en") + # Check english is selected. + expect(select).to_have_value("en") + # Check default jobs search columns. page.goto(f"{BASE_URL}/jobs/search") headers = page.locator("table#search_table thead tr th") diff --git a/clockwork_tools/clockwork_tools/client.py b/clockwork_tools/clockwork_tools/client.py index f436ffb9..b239fe16 100644 --- a/clockwork_tools/clockwork_tools/client.py +++ b/clockwork_tools/clockwork_tools/client.py @@ -60,7 +60,7 @@ def _get_headers(self): encoded_s = str(encoded_bytes, "utf-8") return {"Authorization": f"Basic {encoded_s}"} - def _request(self, endpoint, params, method="GET"): + def _request(self, endpoint, params, method="GET", send_json=True): """Helper method for REST API calls. Internal helper method to make the calls to the REST API endpoints @@ -69,6 +69,9 @@ def _request(self, endpoint, params, method="GET"): Args: endpoint (str): The REST endpoint, omitting the server address. params (dict): Arguments to be provided to the REST endpoint. + send_json (bool): Optional. If True and if method is PUT, + then request will be sent as a JSON request. + Returns: Depends on the call made. @@ -86,9 +89,14 @@ def _request(self, endpoint, params, method="GET"): complete_address, params=params, headers=self._get_headers() ) elif method == "PUT": - response = requests.put( - complete_address, data=params, headers=self._get_headers() - ) + if send_json: + headers = self._get_headers() + headers["Content-type"] = "application/json" + response = requests.put(complete_address, json=params, headers=headers) + else: + response = requests.put( + complete_address, data=params, headers=self._get_headers() + ) # Check code instead and raise exception if it's the wrong one. if response.status_code == 200: @@ -154,6 +162,60 @@ def jobs_one(self, job_id: str = None, cluster_name: str = None) -> dict[str, an params[k] = a return self._request(endpoint, params) + def get_user_props(self, job_id: str, cluster_name: str) -> dict[str, any]: + """REST call to api/v1/clusters/jobs/user_props/get. + + Call for get_user_props(). + + Args: + job_id (str): ID of job to retrieve props + cluster_name (str): Name of job cluster + + Returns: + dict[any,any]: props. + """ + endpoint = "api/v1/clusters/jobs/user_props/get" + params = {"job_id": job_id, "cluster_name": cluster_name} + return self._request(endpoint, params) + + def set_user_props( + self, job_id: str, cluster_name: str, updates: dict + ) -> dict[str, any]: + """REST call to api/v1/clusters/jobs/user_props/set. + + Call for set_user_props(). + + Args: + job_id (str): ID of job to retrieve props + cluster_name (str): Name of job cluster + updates (dict): Dict of props to update. + + Returns: + dict[any,any]: Returns the updated props. + """ + endpoint = "api/v1/clusters/jobs/user_props/set" + params = {"job_id": job_id, "cluster_name": cluster_name, "updates": updates} + return self._request(endpoint, params, method="PUT") + + def delete_user_props( + self, job_id: str, cluster_name: str, keys: str | list + ) -> dict[str, any]: + """REST call to api/v1/clusters/jobs/user_props/delete. + + Call for delete_user_props(). + + Args: + job_id (str): ID of job to retrieve props + cluster_name (str): Name of job cluster + keys (list): List of keys to delete. + + Returns: + dict[any,any]: Returns the updated props. + """ + endpoint = "api/v1/clusters/jobs/user_props/delete" + params = {"job_id": job_id, "cluster_name": cluster_name, "keys": keys} + return self._request(endpoint, params, method="PUT") + def jobs_user_dict_update( self, job_id: str = None, cluster_name: str = None, update_pairs: dict = {} ) -> dict[str, any]: @@ -191,7 +253,7 @@ def jobs_user_dict_update( # Due to current constraints, we have to pass "update_pairs" # as a string representing a structure in json. params["update_pairs"] = json.dumps(update_pairs) - return self._request(endpoint, params, method="PUT") + return self._request(endpoint, params, method="PUT", send_json=False) def nodes_list(self, cluster_name: str = None) -> list[dict[str, any]]: """REST call to api/v1/clusters/nodes/list. diff --git a/clockwork_tools_test/test_mt_job_user_props.py b/clockwork_tools_test/test_mt_job_user_props.py new file mode 100644 index 00000000..13d01bb7 --- /dev/null +++ b/clockwork_tools_test/test_mt_job_user_props.py @@ -0,0 +1,97 @@ +import random + + +def _get_test_user_props(mtclient, fake_data): + # Find an entry that's associated with the user that's currently logged in. + # This becomes the ground truth against which we compare the retrieved user props. + LD_candidates = [ + D_job_user_props_entry + for D_job_user_props_entry in fake_data["job_user_props"] + if ( + D_job_user_props_entry["mila_email_username"] == mtclient.email + and len(D_job_user_props_entry["props"]) > 0 + ) + ] + assert ( + len(LD_candidates) > 0 + ), "There should be at least one job_user_props entry for the user that's currently logged in." + D_job_user_props_entry = random.choice(LD_candidates) + + job_id = D_job_user_props_entry["job_id"] + cluster_name = D_job_user_props_entry["cluster_name"] + original_props = D_job_user_props_entry["props"] + return job_id, cluster_name, original_props + + +def test_cw_tools_get_user_props(mtclient, fake_data): + job_id, cluster_name, original_props = _get_test_user_props(mtclient, fake_data) + props = mtclient.get_user_props(job_id, cluster_name) + assert original_props == props + + +def test_cw_tools_set_and_delete_user_props(mtclient, fake_data): + job_id, cluster_name, original_props = _get_test_user_props(mtclient, fake_data) + + assert original_props + # Test we can replace existing user prop value. + an_existing_prop_name = next(iter(original_props.keys())) + a_new_value = "a totally new value" + # We check new value is not already associated to existing prop name + assert original_props[an_existing_prop_name] != a_new_value + # Set new value + props = mtclient.set_user_props( + job_id, cluster_name, {an_existing_prop_name: a_new_value} + ) + # Check set returns same value as get + assert props == mtclient.get_user_props(job_id, cluster_name) + # Check returned props + assert len(props) == len(original_props) + # Props should be original props, except that `an_existing_prop_name` is now associated to `a_new_value`. + assert props == { + an_existing_prop_name: a_new_value, + **{ + key: value + for key, value in original_props.items() + if key != an_existing_prop_name + }, + } + + # Get back to original props for tests below + props = mtclient.set_user_props(job_id, cluster_name, original_props) + # Check set and get + assert props == mtclient.get_user_props(job_id, cluster_name) == original_props + + # Define new props to add + new_props = {"a new name": "a new prop", "aa": "aa", "bb": "bb", "cc": "cc"} + # We check new props are not already in original props + for new_prop_name in new_props.keys(): + assert new_prop_name not in original_props + + # Test set new props + props = mtclient.set_user_props(job_id, cluster_name, new_props) + assert len(props) == len(original_props) + len(new_props) + assert props == {**original_props, **new_props} + + # Test delete a prop using a str + assert mtclient.delete_user_props(job_id, cluster_name, "aa") == "" + props = mtclient.get_user_props(job_id, cluster_name) + assert len(props) == len(original_props) + len(new_props) - 1 + assert props == { + **original_props, + "a new name": "a new prop", + "bb": "bb", + "cc": "cc", + } + + # Test delete a prop using a list with 1 string + assert mtclient.delete_user_props(job_id, cluster_name, ["a new name"]) == "" + props = mtclient.get_user_props(job_id, cluster_name) + assert len(props) == len(original_props) + len(new_props) - 2 + assert props == {**original_props, "bb": "bb", "cc": "cc"} + + # Test delete multiple props using a list with many strings + assert mtclient.delete_user_props(job_id, cluster_name, ["bb", "cc"]) == "" + props = mtclient.get_user_props(job_id, cluster_name) + assert props == original_props + + # As props are not original_props, no need to clean-up code here! diff --git a/clockwork_web/core/job_user_props_helper.py b/clockwork_web/core/job_user_props_helper.py new file mode 100644 index 00000000..6089678c --- /dev/null +++ b/clockwork_web/core/job_user_props_helper.py @@ -0,0 +1,159 @@ +"""Internal functions to manage job-user props.""" +from ..db import get_db +import json + + +# Max length allowed for a JSON-string representing a job-user props dict. +# Currently, 2 Mb. +MAX_PROPS_LENGTH = 2 * 1024 * 1024 + + +class HugeUserPropsError(ValueError): + """Exception raised when user props are too huge in database.""" + + +def get_user_props(job_id: str, cluster_name: str, mila_email_username: str) -> dict: + """ + Get job-user props. + + Parameters: + job_id ID of job for which we want to get user props. + cluster_name Name of cluster to which the job belongs. + mila_email_username Email of user who sets the props we want to get. + + Returns: + Dictionary of user-props, empty if no props were found. + Each prop is a key-value pair in returned dictionary. + """ + # Get matching MongoDB document. + props_dict = _get_user_props_document(job_id, cluster_name, mila_email_username) + # Return props if available. + return props_dict.get("props", {}) + + +def set_user_props( + job_id: str, cluster_name: str, updates: dict, mila_email_username: str +) -> dict: + """ + Update job-user-props. + + Parameters: + job_id ID of job for which we want to update user props. + cluster_name Name of cluster to which the job belongs. + updates Dictionary of props to add. + Each key-value represents a prop. + mila_email_username Email of user who wants to update his props. + + Returns: + Dictionary of Updated job-user props. + """ + # Get previous props and check that + # previous + updates props do not exceed a size limit. + previous_doc = _get_user_props_document(job_id, cluster_name, mila_email_username) + previous_props = previous_doc.get("props", {}) + new_props = previous_props.copy() + new_props.update(updates) + if _get_dict_size(new_props) > MAX_PROPS_LENGTH: + raise HugeUserPropsError( + f"Too huge job-user props: maximum {_get_megabytes(MAX_PROPS_LENGTH)} Mbytes " + f"(previous: {_get_megabytes(_get_dict_size(previous_props))} Mbytes, " + f"updates: {_get_megabytes(_get_dict_size(updates))} Mbytes)" + ) + + # Update or insert job-user props. + mc = get_db() + db = mc["job_user_props"] + if "_id" in previous_doc: + db.update_one({"_id": previous_doc["_id"]}, {"$set": {"props": new_props}}) + else: + db.insert_one( + { + "job_id": job_id, + "cluster_name": cluster_name, + "mila_email_username": mila_email_username, + "props": new_props, + } + ) + + return new_props + + +def delete_user_props(job_id, cluster_name, key_or_keys, mila_email_username: str): + """ + Delete some job-user props. + + Parameters: + job_id ID of job for which we want to delete some user props. + cluster_name Name of cluster to which the job belongs. + key_or_keys Either a key or a sequence of keys, + representing prop names to delete. + mila_email_username Email of user who wants to delete props. + """ + # Parse key_or_keys to get keys to delete. + if isinstance(key_or_keys, str): + keys = [key_or_keys] + else: + assert isinstance(key_or_keys, (list, tuple, set)) + keys = list(key_or_keys) + + # Get previous props + previous_doc = _get_user_props_document(job_id, cluster_name, mila_email_username) + previous_props = previous_doc.get("props", {}) + new_props = previous_props.copy() + # Remove keys + for key in keys: + new_props.pop(key, None) + # Register in MongoDB + if len(new_props) < len(previous_props): + mc = get_db() + db = mc["job_user_props"] + db.update_one({"_id": previous_doc["_id"]}, {"$set": {"props": new_props}}) + + +def _get_user_props_document( + job_id: str, cluster_name: str, mila_email_username: str +) -> dict: + """ + Get MongoDB document representing job-user props corresponding to given parameters. + + Parameters: + job_id ID of job for which we want to get user props. + cluster_name Name of cluster to which the job belongs. + mila_email_username Email of user who sets the props we want to get. + + Returns: + MongoDB document representing job-user props. + Document is a dictionary with following format: + { + job_id: str # ID of job associated to props. + cluster_name: str # cluster name of job associated to props. + mila_email_username: str # user who created these props. + props: dict # actual user-props, each prop is a key-value. + } + """ + mc = get_db() + result = list( + mc["job_user_props"].find( + { + "job_id": job_id, + "cluster_name": cluster_name, + "mila_email_username": mila_email_username, + } + ) + ) + if not result: + return {} + else: + assert len(result) == 1 + (props,) = result + return props + + +def _get_dict_size(dct: dict) -> int: + """Get length of JSON-string representation for given dictionary.""" + return len(json.dumps(dct)) + + +def _get_megabytes(size: int) -> float: + """Convert bytes to megabytes.""" + return size / (1024 * 1024) diff --git a/clockwork_web/core/jobs_helper.py b/clockwork_web/core/jobs_helper.py index 3ee9e8fe..50dbd6a7 100644 --- a/clockwork_web/core/jobs_helper.py +++ b/clockwork_web/core/jobs_helper.py @@ -167,9 +167,7 @@ def get_filtered_and_paginated_jobs( mc["job_user_props"].find( combine_all_mongodb_filters( { - "job_id": { - "$in": [int(job["slurm"]["job_id"]) for job in LD_jobs] - }, + "job_id": {"$in": [job["slurm"]["job_id"] for job in LD_jobs]}, "mila_email_username": current_user.mila_email_username, } ) @@ -190,7 +188,7 @@ def get_filtered_and_paginated_jobs( for job in LD_jobs: key = ( current_user.mila_email_username, - int(job["slurm"]["job_id"]), + job["slurm"]["job_id"], job["slurm"]["cluster_name"], ) if key in user_props_map: diff --git a/clockwork_web/rest_routes/jobs.py b/clockwork_web/rest_routes/jobs.py index 382fc68f..9b45e05c 100644 --- a/clockwork_web/rest_routes/jobs.py +++ b/clockwork_web/rest_routes/jobs.py @@ -23,6 +23,12 @@ get_jobs, ) from clockwork_web.core.utils import to_boolean, get_custom_array_from_request_args +from clockwork_web.core.job_user_props_helper import ( + get_user_props, + set_user_props, + delete_user_props, + HugeUserPropsError, +) from flask import Blueprint @@ -126,6 +132,136 @@ def route_api_v1_jobs_one(): return jsonify(D_job) +@flask_api.route("/jobs/user_props/get") +@authentication_required +def route_user_props_get(): + """ + Endpoint to get user props. + + Parameters: job_id (str), cluster_name (str) + + Return: user props + """ + current_user_id = g.current_user_with_rest_auth["mila_email_username"] + logging.info( + f"clockwork REST route: /jobs/user_props/get - current_user_with_rest_auth={current_user_id}" + ) + + # Retrieve the requested job ID + job_id = request.values.get("job_id", None) + if job_id is None: + return jsonify("Missing argument job_id."), 400 # bad request + + # Retrieve the requested cluster name + cluster_name = request.values.get("cluster_name", None) + if cluster_name is None: + return jsonify("Missing argument cluster_name."), 400 # bad request + + # Get props, using current_user_id as mila email username. + props = get_user_props(job_id, cluster_name, current_user_id) + return jsonify(props) + + +@flask_api.route("/jobs/user_props/set", methods=["PUT"]) +@authentication_required +def route_user_props_set(): + """ + Endpoint to set user props + + Parameters: job_id (str), cluster_name (str), updates (dict) + + Return: updated user props + """ + current_user_id = g.current_user_with_rest_auth["mila_email_username"] + logging.info( + f"clockwork REST route: /jobs/user_props/set - current_user_with_rest_auth={current_user_id}" + ) + + if not request.is_json: + return jsonify("Expected a JSON request"), 400 # bad request + + params = request.get_json() + + # Retrieve the requested job ID + job_id = params.get("job_id", None) + if job_id is None: + return jsonify("Missing argument job_id."), 400 # bad request + + # Retrieve the requested cluster name + cluster_name = params.get("cluster_name", None) + if cluster_name is None: + return jsonify("Missing argument cluster_name."), 400 # bad request + + # Retrieve the requested updates. + updates = params.get("updates", None) + if updates is None: + return jsonify(f"Missing argument 'updates'."), 500 + if not isinstance(updates, dict): + return ( + jsonify( + f"Expected `updates` to be a dict, but instead it is of type {type(updates)}: {updates}." + ), + 500, + ) + + # Set props, using current_user_id as mila email username. + try: + props = set_user_props(job_id, cluster_name, updates, current_user_id) + return jsonify(props) + except HugeUserPropsError as inst: + # If props size limit error occurs, return it as an HTTP 500 error. + return jsonify(str(inst)), 500 + except: + return jsonify("Failed to set user props."), 500 + + +@flask_api.route("/jobs/user_props/delete", methods=["PUT"]) +@authentication_required +def route_user_props_delete(): + """ + Endpoint to delete user props. + + Parameters: job_id (str), cluster_name (str), keys (string or list of strings) + + Return: updated user props + """ + current_user_id = g.current_user_with_rest_auth["mila_email_username"] + logging.info( + f"clockwork REST route: /jobs/user_props/delete - current_user_with_rest_auth={current_user_id}" + ) + + if not request.is_json: + return jsonify("Expected a JSON request"), 400 # bad request + + params = request.get_json() + + # Retrieve the requested job ID + job_id = params.get("job_id", None) + if job_id is None: + return jsonify("Missing argument job_id."), 400 # bad request + + # Retrieve the requested cluster name + cluster_name = params.get("cluster_name", None) + if cluster_name is None: + return jsonify("Missing argument cluster_name."), 400 # bad request + + # Retrieve the requested keys. + keys = params.get("keys", None) + if keys is None: + return jsonify(f"Missing argument 'keys'."), 500 + if not isinstance(keys, (str, list)): + return ( + jsonify( + f"Expected `keys` to be a string or list of strings, instead it is of type {type(keys)}: {keys}." + ), + 500, + ) + + # Delete props, using current_user_id as mila email username. + delete_user_props(job_id, cluster_name, keys, current_user_id) + return jsonify("") + + # Note that this whole `user_dict_update` thing needs to be rewritten # in order to use Olivier's proposal about jobs properties # being visible only to the users that set them, diff --git a/clockwork_web_test/test_rest_job_user_props.py b/clockwork_web_test/test_rest_job_user_props.py new file mode 100644 index 00000000..3f115379 --- /dev/null +++ b/clockwork_web_test/test_rest_job_user_props.py @@ -0,0 +1,256 @@ +import random +from clockwork_web.config import get_config +from clockwork_web.core.job_user_props_helper import MAX_PROPS_LENGTH + + +def _get_test_user_props(fake_data): + email = get_config("clockwork.test.email") + # Find an entry that's associated with the user that's currently logged in. + # This becomes the ground truth against which we compare the retrieved user props. + LD_candidates = [ + D_job_user_props_entry + for D_job_user_props_entry in fake_data["job_user_props"] + if ( + D_job_user_props_entry["mila_email_username"] == email + and len(D_job_user_props_entry["props"]) > 0 + ) + ] + assert ( + len(LD_candidates) > 0 + ), "There should be at least one job_user_props entry for the user that's currently logged in." + D_job_user_props_entry = random.choice(LD_candidates) + + job_id = D_job_user_props_entry["job_id"] + cluster_name = D_job_user_props_entry["cluster_name"] + original_props = D_job_user_props_entry["props"] + return job_id, cluster_name, original_props + + +def test_jobs_user_props_get(client, valid_rest_auth_headers, fake_data): + job_id, cluster_name, original_props = _get_test_user_props(fake_data) + response = client.get( + f"/api/v1/clusters/jobs/user_props/get?cluster_name={cluster_name}&job_id={job_id}", + headers=valid_rest_auth_headers, + ) + assert response.content_type == "application/json" + assert response.status_code == 200 + props = response.get_json() + assert props == original_props + + +def test_jobs_user_props_get_unknown_cluster( + client, valid_rest_auth_headers, fake_data +): + job_id, _, _ = _get_test_user_props(fake_data) + cluster_name = "unknown_cluster" + response = client.get( + f"/api/v1/clusters/jobs/user_props/get?cluster_name={cluster_name}&job_id={job_id}", + headers=valid_rest_auth_headers, + ) + assert response.content_type == "application/json" + assert response.status_code == 200 + props = response.get_json() + assert props == {} + + +def test_jobs_user_props_get_unknown_job_id(client, valid_rest_auth_headers, fake_data): + _, cluster_name, _ = _get_test_user_props(fake_data) + job_id = "unknown_job_id" + response = client.get( + f"/api/v1/clusters/jobs/user_props/get?cluster_name={cluster_name}&job_id={job_id}", + headers=valid_rest_auth_headers, + ) + assert response.content_type == "application/json" + assert response.status_code == 200 + props = response.get_json() + assert props == {} + + +def test_jobs_user_props_set(client, valid_rest_auth_headers, fake_data): + job_id, cluster_name, original_props = _get_test_user_props(fake_data) + assert "other name" not in original_props + assert "other name 2" not in original_props + response = client.put( + f"/api/v1/clusters/jobs/user_props/set", + json={ + "job_id": job_id, + "cluster_name": cluster_name, + "updates": {"other name": "other value", "other name 2": "other value 2"}, + }, + headers=valid_rest_auth_headers, + ) + assert response.content_type == "application/json" + assert response.status_code == 200 + props = response.get_json() + assert len(props) == len(original_props) + 2 + assert props == { + **original_props, + "other name": "other value", + "other name 2": "other value 2", + } + + # Change value for existing prop 'other name` + response = client.put( + f"/api/v1/clusters/jobs/user_props/set", + json={ + "job_id": job_id, + "cluster_name": cluster_name, + "updates": {"other name": "new value for prop 'other name'"}, + }, + headers=valid_rest_auth_headers, + ) + assert response.content_type == "application/json" + assert response.status_code == 200 + props = response.get_json() + assert len(props) == len(original_props) + 2 + assert props == { + **original_props, + "other name": "new value for prop 'other name'", + "other name 2": "other value 2", + } + + # Back to default props + client.put( + f"/api/v1/clusters/jobs/user_props/delete", + json={ + "job_id": job_id, + "cluster_name": cluster_name, + "keys": ["other name", "other name 2"], + }, + headers=valid_rest_auth_headers, + ) + assert ( + client.get( + f"/api/v1/clusters/jobs/user_props/get?cluster_name={cluster_name}&job_id={job_id}", + headers=valid_rest_auth_headers, + ).get_json() + == original_props + ) + + +def test_jobs_user_props_delete(client, valid_rest_auth_headers, fake_data): + # Set some props. + job_id, cluster_name, original_props = _get_test_user_props(fake_data) + assert "other name" not in original_props + assert "dino" not in original_props + response = client.put( + f"/api/v1/clusters/jobs/user_props/set", + json={ + "job_id": job_id, + "cluster_name": cluster_name, + "updates": { + "other name": "other value", + "dino": "saurus", + "aa": "aa", + "bb": "bb", + "cc": "cc", + }, + }, + headers=valid_rest_auth_headers, + ) + assert response.content_type == "application/json" + assert response.status_code == 200 + props = response.get_json() + assert len(props) == len(original_props) + 5 + assert props == { + **original_props, + "other name": "other value", + "dino": "saurus", + "aa": "aa", + "bb": "bb", + "cc": "cc", + } + + # Delete a prop with keys as a list of 1 string. + response = client.put( + f"/api/v1/clusters/jobs/user_props/delete", + json={ + "job_id": job_id, + "cluster_name": cluster_name, + "keys": ["dino"], + }, + headers=valid_rest_auth_headers, + ) + assert response.content_type == "application/json" + assert response.status_code == 200 + assert response.get_json() == "" + + response = client.get( + f"/api/v1/clusters/jobs/user_props/get?cluster_name={cluster_name}&job_id={job_id}", + headers=valid_rest_auth_headers, + ) + assert response.status_code == 200 + props = response.get_json() + assert len(props) == len(original_props) + 4 + assert props == { + **original_props, + "other name": "other value", + "aa": "aa", + "bb": "bb", + "cc": "cc", + } + + # Delete some props with keys as a list of strings, including 1 unknown prop. + response = client.put( + f"/api/v1/clusters/jobs/user_props/delete", + json={ + "job_id": job_id, + "cluster_name": cluster_name, + "keys": ["aa", "bb", "unknown prop", "cc"], + }, + headers=valid_rest_auth_headers, + ) + assert response.content_type == "application/json" + assert response.status_code == 200 + assert response.get_json() == "" + + response = client.get( + f"/api/v1/clusters/jobs/user_props/get?cluster_name={cluster_name}&job_id={job_id}", + headers=valid_rest_auth_headers, + ) + assert response.status_code == 200 + props = response.get_json() + assert len(props) == len(original_props) + 1 + assert props == {**original_props, "other name": "other value"} + + # Back to default props + # And delete a prop with keys as a string + client.put( + f"/api/v1/clusters/jobs/user_props/delete", + json={"job_id": job_id, "cluster_name": cluster_name, "keys": "other name"}, + headers=valid_rest_auth_headers, + ) + assert ( + client.get( + f"/api/v1/clusters/jobs/user_props/get?cluster_name={cluster_name}&job_id={job_id}", + headers=valid_rest_auth_headers, + ).get_json() + == original_props + ) + + +def test_size_limit_for_jobs_user_props_set(client, valid_rest_auth_headers, fake_data): + job_id, cluster_name, original_props = _get_test_user_props(fake_data) + assert "other name" not in original_props + huge_text = "x" * (MAX_PROPS_LENGTH + 1) + response = client.put( + f"/api/v1/clusters/jobs/user_props/set", + json={ + "job_id": job_id, + "cluster_name": cluster_name, + "updates": {"other name": huge_text}, + }, + headers=valid_rest_auth_headers, + ) + assert response.content_type == "application/json" + assert response.status_code == 500 + assert response.get_json().startswith("Too huge job-user props: maximum 2.0 Mbytes") + + # Props should have not changed. + response = client.get( + f"/api/v1/clusters/jobs/user_props/get?cluster_name={cluster_name}&job_id={job_id}", + headers=valid_rest_auth_headers, + ) + assert response.content_type == "application/json" + assert response.status_code == 200 + assert response.get_json() == original_props diff --git a/dev.sh b/dev.sh index 4e6331f4..fd105fc8 100644 --- a/dev.sh +++ b/dev.sh @@ -7,6 +7,6 @@ docker build -t clockwork_dev -f clockwork_dev.Dockerfile . . ./env.sh # Make sure to cleanup on exit -trap "docker-compose down && docker-compose rm -fv" EXIT +trap "docker compose down && docker compose rm -fv" EXIT -docker-compose run --service-ports clockwork_dev +docker compose run --service-ports clockwork_dev diff --git a/produce_fake_data.sh b/produce_fake_data.sh index 25127f09..fb2e1317 100644 --- a/produce_fake_data.sh +++ b/produce_fake_data.sh @@ -5,7 +5,7 @@ set -eu . ./env.sh # This is to ensure that there aren't lingering containers after the test script exits. -trap "docker-compose down && docker-compose rm -fv" EXIT +trap "docker compose down && docker compose rm -fv" EXIT # Run the Docker containers -docker-compose run clockwork_scripts +docker compose run clockwork_scripts diff --git a/scripts/store_huge_fake_data_in_db.py b/scripts/store_huge_fake_data_in_db.py deleted file mode 100644 index b654eb2e..00000000 --- a/scripts/store_huge_fake_data_in_db.py +++ /dev/null @@ -1,271 +0,0 @@ -""" -This script is inspired from `store_fake_data_in_db.py` for same usage, i.e. with "dev.sh". -The difference is that here, we can call this script with parameters to control -how much data will be inserted in database. This allows to test database -when populated with a huge amount of data. -""" - -import argparse -import os.path -import sys -import json - -from clockwork_web.config import register_config -from slurm_state.mongo_client import get_mongo_client -from slurm_state.config import get_config - - -DEFAULT_NB_JOBS = 1_000_000 -DEFAULT_NB_DICTS = DEFAULT_NB_JOBS -DEFAULT_NB_PROPS_PER_DICT = 4 - - -def main(argv): - # Retrieve the arguments passed to the script - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter - ) - group = parser.add_mutually_exclusive_group() - group.add_argument( - "-j", - "--nb-student-jobs", - action="append", - type=str, - help=( - "Number of job for a specific student, in format: =. " - "Accept multiple declarations. Example: -j student00=100 -j student05=1900. " - " must be a valid student available in fake_data.json's users. " - " is just the student name, without '@mila.quebec'" - ), - ) - group.add_argument( - "--nb-jobs", - type=int, - default=DEFAULT_NB_JOBS, - help="Number of jobs to add. May be 0 (no job added). Mutually exclusive with --nb-student-jobs.", - ) - parser.add_argument( - "--nb-dicts", - type=int, - default=DEFAULT_NB_DICTS, - help="Number of job-user dicts to add. May be 0 (no job added).", - ) - parser.add_argument( - "--nb-props-per-dict", - type=int, - default=DEFAULT_NB_PROPS_PER_DICT, - help=f"Number of key-value pairs in each job-user dict.", - ) - parser.add_argument( - "--props-username", - type=str, - default="student00@mila.quebec", - help="Email of user who creates job-user dicts.", - ) - parser.add_argument( - "--disable-index", - action="store_true", - help="If specified, will not create MongoDB index.", - ) - args = parser.parse_args(argv[1:]) - print(args) - - # Register the elements to access the database - register_config("mongo.connection_string", "") - register_config("mongo.database_name", "clockwork") - - # Store the generated fake data in the database - store_data_in_db( - nb_jobs=args.nb_jobs, - nb_student_jobs=args.nb_student_jobs, - nb_dicts=args.nb_dicts, - nb_props_per_dict=args.nb_props_per_dict, - props_username=args.props_username, - disable_index=args.disable_index, - ) - - -def store_data_in_db(**kwargs): - # Open the database and insert the contents. - client = get_mongo_client() - populate_fake_data(client[get_config("mongo.database_name")], **kwargs) - - -def populate_fake_data(db_insertion_point, **kwargs): - disable_index = kwargs.pop("disable_index", False) - - print("Generating huge fake data") - E = _generate_huge_fake_data(**kwargs) - print("Generated huge fake data") - - # Drop any collection (and related index) before. - for k in ["users", "jobs", "nodes", "gpu", "job_user_props"]: - db_insertion_point[k].drop() - # This should verify we do not have collection indexes. - assert not list(db_insertion_point[k].list_indexes()) - - if not disable_index: - print("Generate MongoDB index.") - # Create indices. This isn't half as important as when we're - # dealing with large quantities of data, but it's part of the - # set up for the database. - db_insertion_point["jobs"].create_index( - [ - ("slurm.job_id", 1), - ("slurm.cluster_name", 1), - ("cw.mila_email_username", 1), - ], - name="job_id_and_cluster_name", - ) - db_insertion_point["nodes"].create_index( - [("slurm.name", 1), ("slurm.cluster_name", 1)], - name="name_and_cluster_name", - ) - db_insertion_point["users"].create_index( - [("mila_email_username", 1)], name="users_email_index" - ) - db_insertion_point["gpu"].create_index([("name", 1)], name="gpu_name") - db_insertion_point["job_user_props"].create_index( - [ - ("mila_email_username", 1), - ("job_id", 1), - ("cluster_name", 1), - ], - name="job_user_props_index", - ) - - # This should verify we do have collection indexes. - for k in ["users", "jobs", "nodes", "gpu", "job_user_props"]: - assert list(db_insertion_point[k].list_indexes()) - - for k in ["users", "jobs", "nodes", "gpu", "job_user_props"]: - # Anyway clean before inserting - db_insertion_point[k].delete_many({}) - if k in E and E[k]: - print(f"Inserting {k}, {len(E[k])} value(s)") - db_insertion_point[k].insert_many(E[k]) - # Check count - assert db_insertion_point[k].count_documents({}) == len(E[k]) - print("Inserted", k) - - -def _generate_huge_fake_data( - nb_jobs=DEFAULT_NB_JOBS, - nb_student_jobs=None, - nb_dicts=DEFAULT_NB_DICTS, - nb_props_per_dict=DEFAULT_NB_PROPS_PER_DICT, - props_username="student00@mila.quebec", -): - student_to_nb_jobs = [] - if nb_student_jobs is not None: - for desc in nb_student_jobs: - student_name, str_nb_student_jobs = desc.split("=") - nb_student_jobs = int(str_nb_student_jobs.strip()) - student_to_nb_jobs.append((student_name.strip(), nb_student_jobs)) - else: - assert nb_jobs >= 0 - - # Get users from fake_data.json - json_file = os.path.join( - os.path.dirname(os.path.abspath(__file__)), - "..", - "test_common", - "fake_data.json", - ) - assert os.path.exists( - json_file - ), f"Failed to find the fake data file to get fake users for testing: {json_file}." - with open(json_file) as f: - E = json.load(f) - users = E["users"] - assert len(users) == 20 - - jobs = [] - - # populate jobs - if student_to_nb_jobs: - user_map = {user["mila_email_username"]: user for user in users} - assert len(user_map) == len(users) - job_id = 0 - for student_name, nb_student_jobs in student_to_nb_jobs: - student_email = f"{student_name}@mila.quebec" - user = user_map[student_email] - for i in range(nb_student_jobs): - job_id += 1 - jobs.append( - _gen_new_job( - job_id, user["cc_account_username"], user["mila_email_username"] - ) - ) - print(f"Student {student_email}: {nb_student_jobs} jobs") - assert job_id == len(jobs) - else: - for i in range(nb_jobs): - # Pick a user - user = users[i % len(users)] - # Then create job - job_id = i + 1 - jobs.append( - _gen_new_job( - job_id, user["cc_account_username"], user["mila_email_username"] - ) - ) - - # populate job-user-dicts - job_user_dicts = [ - { - "mila_email_username": props_username, - "job_id": i + 1, - "cluster_name": "beluga", - "props": { - f"prop_{j + 1}_for_job_{i + 1}": f"I am user dict prop {j + 1} for job ID {i + 1}" - for j in range(nb_props_per_dict) - }, - } - for i in range(nb_dicts) - ] - - print( - f"Jobs: {len(jobs)}, dicts: {len(job_user_dicts)}, props per dict: {nb_props_per_dict}" - ) - return {"users": users, "jobs": jobs, "job_user_props": job_user_dicts} - - -def _gen_new_job(job_id, slurm_username, mila_email_username): - return { - "slurm": { - "account": "def-patate-rrg", - "cluster_name": "beluga", - "time_limit": 4320, - "submit_time": 1681680327, - "start_time": 0, - "end_time": 0, - "exit_code": "SUCCESS:0", - "array_job_id": "0", - "array_task_id": "None", - "job_id": str(job_id), - "name": f"job_name_{job_id}", - "nodes": "None assigned", - "partition": "other_fun_partition", - "job_state": "PENDING", - "tres_allocated": {}, - "tres_requested": { - "num_cpus": 80, - "mem": 95000, - "num_nodes": 1, - "billing": 80, - }, - "username": slurm_username, - "working_directory": "/a809/b333/c569", - }, - "cw": { - "mila_email_username": mila_email_username, - "last_slurm_update": 1686248596.476063, - "last_slurm_update_by_sacct": 1686248596.476063, - }, - "user": {}, - } - - -if __name__ == "__main__": - main(sys.argv) diff --git a/test.sh b/test.sh index 166a2d06..bc76c19e 100644 --- a/test.sh +++ b/test.sh @@ -10,9 +10,9 @@ docker build -t scripts_test -f scripts_test/Dockerfile . . ./env.sh # This is to ensure that there aren't lingering containers after the test script exits. -trap "docker-compose down && docker-compose rm -fv" EXIT +trap "docker compose down && docker compose rm -fv" EXIT -docker-compose run clockwork_web_test -docker-compose run clockwork_tools_test -docker-compose run slurm_state_test -docker-compose run scripts_test +docker compose run clockwork_web_test +docker compose run clockwork_tools_test +docker compose run slurm_state_test +docker compose run scripts_test diff --git a/test_common/fake_data.json b/test_common/fake_data.json index 29c18e58..2670cd4f 100644 --- a/test_common/fake_data.json +++ b/test_common/fake_data.json @@ -5965,32 +5965,32 @@ ], "job_user_props": [ { - "mila_email_username": "student00@mila.quebec", - "job_id": 795002, + "mila_email_username": "student01@mila.quebec", + "job_id": "795002", "cluster_name": "mila", "props": { "name": "je suis une user prop 1" } }, { - "mila_email_username": "student00@mila.quebec", - "job_id": 606872, + "mila_email_username": "student01@mila.quebec", + "job_id": "606872", "cluster_name": "mila", "props": { "name": "je suis une user prop 2" } }, { - "mila_email_username": "student00@mila.quebec", - "job_id": 834395, + "mila_email_username": "student01@mila.quebec", + "job_id": "834395", "cluster_name": "graham", "props": { "name": "je suis une user prop 3" } }, { - "mila_email_username": "student00@mila.quebec", - "job_id": 154325, + "mila_email_username": "student01@mila.quebec", + "job_id": "154325", "cluster_name": "graham", "props": { "name": "je suis une user prop 3", @@ -5998,8 +5998,8 @@ } }, { - "mila_email_username": "student00@mila.quebec", - "job_id": 613024, + "mila_email_username": "student01@mila.quebec", + "job_id": "613024", "cluster_name": "graham", "props": { "name": "je suis une user prop 1"