diff --git a/clockwork_frontend_test/test_jobs_search.py b/clockwork_frontend_test/test_jobs_search.py index b109e86b..db5cd55e 100644 --- a/clockwork_frontend_test/test_jobs_search.py +++ b/clockwork_frontend_test/test_jobs_search.py @@ -640,6 +640,59 @@ def test_filter_by_job_array(page: Page): _check_jobs_table(page, JOBS_SEARCH_DEFAULT_TABLE) +def test_filter_by_job_user_props(page: Page): + # Login + page.goto(f"{BASE_URL}/login/testing?user_id=student00@mila.quebec") + # Go to settings. + page.goto(f"{BASE_URL}/settings/") + radio_job_user_props = page.locator("input#jobs_list_job_user_props_toggle") + expect(radio_job_user_props).to_be_checked(checked=False) + # Check column job_user_props. + radio_job_user_props.click() + expect(radio_job_user_props).to_be_checked(checked=True) + # Back to jobs/search. + page.goto(f"{BASE_URL}/jobs/search") + + job_id = page.get_by_text("795002") + expect(job_id).to_have_count(1) + parent_row = page.locator("table#search_table tbody tr").filter(has=job_id) + expect(parent_row).to_have_count(1) + cols = parent_row.locator("td") + link_job_user_prop = cols.nth(4).locator("a") + expect(link_job_user_prop).to_have_count(1) + expect(link_job_user_prop).to_contain_text("name je suis une user prop 1") + link_job_user_prop.click() + expect(page).to_have_url( + f"{BASE_URL}/jobs/search?" + f"user_prop_name=name" + f"&user_prop_content=je+suis+une+user+prop+1" + f"&page_num=1" + ) + _check_jobs_table( + page, + [ + ["mila", "student06 @mila.quebec", "795002"], + ["graham", "student12 @mila.quebec", "613024"], + ], + ) + + filter_reset = page.get_by_title("Reset filter by job user prop") + expect(filter_reset).to_contain_text('User prop name: "je suis une user prop 1"') + filter_reset.click() + + expect(page).to_have_url( + f"{BASE_URL}/jobs/search?user_prop_name=&user_prop_content=&page_num=1" + ) + _check_jobs_table(page, JOBS_SEARCH_DEFAULT_TABLE) + + # Back to default settings. + page.goto(f"{BASE_URL}/settings/") + radio_job_user_props = page.locator("input#jobs_list_job_user_props_toggle") + expect(radio_job_user_props).to_be_checked(checked=True) + radio_job_user_props.click() + expect(radio_job_user_props).to_be_checked(checked=False) + + def test_jobs_table_sorting_by_cluster(page: Page): _load_jobs_search_page(page) expected_content = [ diff --git a/clockwork_frontend_test/test_settings.py b/clockwork_frontend_test/test_settings.py index 2fcd70c3..39d58045 100644 --- a/clockwork_frontend_test/test_settings.py +++ b/clockwork_frontend_test/test_settings.py @@ -223,3 +223,66 @@ def test_jobs_search_columns(page: Page): expect(headers.nth(7)).to_contain_text("Start time") expect(headers.nth(8)).to_contain_text("End time") expect(headers.nth(9)).to_contain_text("Links") + + +def test_jobs_search_column_job_user_props(page: Page): + # Login + page.goto(f"{BASE_URL}/login/testing?user_id=student00@mila.quebec") + # Check default jobs search columns. + page.goto(f"{BASE_URL}/jobs/search") + headers = page.locator("table#search_table thead tr th") + expect(headers).to_have_count(10) + expect(headers.nth(0)).to_contain_text("Cluster") + expect(headers.nth(1)).to_contain_text("User (@mila.quebec)") + expect(headers.nth(2)).to_contain_text("Job ID") + expect(headers.nth(3)).to_contain_text("Job array") + expect(headers.nth(4)).to_contain_text("Job name [:20]") + expect(headers.nth(5)).to_contain_text("Job state") + expect(headers.nth(6)).to_contain_text("Submit time") + expect(headers.nth(7)).to_contain_text("Start time") + expect(headers.nth(8)).to_contain_text("End time") + expect(headers.nth(9)).to_contain_text("Links") + + # Go to settings. + page.goto(f"{BASE_URL}/settings/") + radio_job_user_props = page.locator("input#jobs_list_job_user_props_toggle") + expect(radio_job_user_props).to_be_checked(checked=False) + # Check column job_user_props. + radio_job_user_props.click() + expect(radio_job_user_props).to_be_checked(checked=True) + + # Check column job_user_props is indeed now displayed in jobs/search. + page.goto(f"{BASE_URL}/jobs/search") + headers = page.locator("table#search_table thead tr th") + expect(headers).to_have_count(11) + expect(headers.nth(0)).to_contain_text("Cluster") + expect(headers.nth(1)).to_contain_text("User (@mila.quebec)") + expect(headers.nth(2)).to_contain_text("Job ID") + expect(headers.nth(3)).to_contain_text("Job array") + expect(headers.nth(4)).to_contain_text("Job-user props") + expect(headers.nth(5)).to_contain_text("Job name [:20]") + expect(headers.nth(6)).to_contain_text("Job state") + expect(headers.nth(7)).to_contain_text("Submit time") + expect(headers.nth(8)).to_contain_text("Start time") + expect(headers.nth(9)).to_contain_text("End time") + expect(headers.nth(10)).to_contain_text("Links") + + # Back to default settings. + page.goto(f"{BASE_URL}/settings/") + radio_job_user_props = page.locator("input#jobs_list_job_user_props_toggle") + expect(radio_job_user_props).to_be_checked(checked=True) + radio_job_user_props.click() + expect(radio_job_user_props).to_be_checked(checked=False) + page.goto(f"{BASE_URL}/jobs/search") + headers = page.locator("table#search_table thead tr th") + expect(headers).to_have_count(10) + expect(headers.nth(0)).to_contain_text("Cluster") + expect(headers.nth(1)).to_contain_text("User (@mila.quebec)") + expect(headers.nth(2)).to_contain_text("Job ID") + expect(headers.nth(3)).to_contain_text("Job array") + expect(headers.nth(4)).to_contain_text("Job name [:20]") + expect(headers.nth(5)).to_contain_text("Job state") + expect(headers.nth(6)).to_contain_text("Submit time") + expect(headers.nth(7)).to_contain_text("Start time") + expect(headers.nth(8)).to_contain_text("End time") + expect(headers.nth(9)).to_contain_text("Links") diff --git a/clockwork_web/browser_routes/jobs.py b/clockwork_web/browser_routes/jobs.py index 289c751a..72bbfc18 100644 --- a/clockwork_web/browser_routes/jobs.py +++ b/clockwork_web/browser_routes/jobs.py @@ -101,6 +101,8 @@ def route_search(): - "sort_asc" is an optional integer and used to specify if sorting is ascending (1) or descending (-1). Default is 1. - "job_array" is optional and used to specify the job array in which we are looking for jobs + - "user_prop_name" is optional and used to specify the user prop name associated to jobs we are looking for + - "user_prop_content" is optional and used to specify the user prop value associated to jobs we are looking for .. :quickref: list all Slurm job as formatted html """ @@ -164,6 +166,8 @@ def route_search(): "sort_by": query.sort_by, "sort_asc": query.sort_asc, "job_array": query.job_array, + "user_prop_name": query.user_prop_name, + "user_prop_content": query.user_prop_content, }, ) diff --git a/clockwork_web/core/jobs_helper.py b/clockwork_web/core/jobs_helper.py index a9ef4c72..3ee9e8fe 100644 --- a/clockwork_web/core/jobs_helper.py +++ b/clockwork_web/core/jobs_helper.py @@ -7,6 +7,7 @@ import time from flask.globals import current_app +from flask_login import current_user from ..db import get_db @@ -157,6 +158,44 @@ def get_filtered_and_paginated_jobs( # on the server because not enough memory was allocated to perform the sorting. LD_jobs = list(mc["jobs"].find(mongodb_filter)) + # Get job user props + if LD_jobs and current_user: + user_props_map = {} + # Collect all job user props related to found jobs, + # and store them in a dict with keys (mila email username, job ID, cluster_name) + for user_props in list( + mc["job_user_props"].find( + combine_all_mongodb_filters( + { + "job_id": { + "$in": [int(job["slurm"]["job_id"]) for job in LD_jobs] + }, + "mila_email_username": current_user.mila_email_username, + } + ) + ) + ): + key = ( + user_props["mila_email_username"], + user_props["job_id"], + user_props["cluster_name"], + ) + assert key not in user_props_map + user_props_map[key] = user_props["props"] + + if user_props_map: + # Populate jobs with user props using + # current user email, job ID and job cluster name + # to find related user props in props map. + for job in LD_jobs: + key = ( + current_user.mila_email_username, + int(job["slurm"]["job_id"]), + job["slurm"]["cluster_name"], + ) + if key in user_props_map: + job["job_user_props"] = user_props_map[key] + # Set nbr_total_jobs if want_count: # Get the number of filtered jobs (not paginated) @@ -235,6 +274,8 @@ def get_jobs( sort_by="submit_time", sort_asc=-1, job_array=None, + user_prop_name=None, + user_prop_content=None, ): """ Set up the filters according to the parameters and retrieve the requested jobs from the database. @@ -252,6 +293,8 @@ def get_jobs( sort_asc Whether or not to sort in ascending order (1) or descending order (-1). job_array ID of job array in which we look for jobs. + user_prop_name name of user prop (string) we must find in jobs to look for. + user_prop_content content of user prop (string) we must find in jobs to look for. Returns: A tuple containing: @@ -259,6 +302,24 @@ def get_jobs( - the total number of jobs corresponding of the filters in the databse, if want_count has been set to True, None otherwise, as second element """ + # If job user prop is specified, + # get job indices from jobs associated to this prop. + if user_prop_name is not None and user_prop_content is not None: + mc = get_db() + props_job_ids = [ + str(user_props["job_id"]) + for user_props in mc["job_user_props"].find( + combine_all_mongodb_filters( + {f"props.{user_prop_name}": user_prop_content} + ) + ) + ] + if job_ids: + # If job ids where provided, make intersection between given job ids and props job ids. + job_ids = list(set(props_job_ids) & set(job_ids)) + else: + # Otherwise, just use props job ids. + job_ids = props_job_ids # Set up and combine filters filter = get_global_filter( @@ -405,6 +466,7 @@ def get_jobs_properties_list_per_page(): "user", "job_id", "job_array", + "job_user_props", "job_name", "job_state", "start_time", diff --git a/clockwork_web/core/search_helper.py b/clockwork_web/core/search_helper.py index 8d80b33e..2650c201 100644 --- a/clockwork_web/core/search_helper.py +++ b/clockwork_web/core/search_helper.py @@ -21,6 +21,8 @@ def parse_search_request(user, args, force_pagination=True): want_count = to_boolean(want_count) job_array = args.get("job_array", type=int, default=None) + user_prop_name = args.get("user_prop_name", type=str, default=None) or None + user_prop_content = args.get("user_prop_content", type=str, default=None) or None default_page_number = "1" if force_pagination else None @@ -71,6 +73,8 @@ def parse_search_request(user, args, force_pagination=True): sort_asc=sort_asc, want_count=want_count, job_array=job_array, + user_prop_name=user_prop_name, + user_prop_content=user_prop_content, ) ######################### @@ -115,5 +119,7 @@ def search_request(user, args, force_pagination=True): sort_by=query.sort_by, sort_asc=query.sort_asc, job_array=query.job_array, + user_prop_name=query.user_prop_name, + user_prop_content=query.user_prop_content, ) return (query, jobs, nbr_total_jobs) diff --git a/clockwork_web/core/users_helper.py b/clockwork_web/core/users_helper.py index 47b5fd74..05a519f4 100644 --- a/clockwork_web/core/users_helper.py +++ b/clockwork_web/core/users_helper.py @@ -605,7 +605,6 @@ def render_template_with_user_settings(template_name_or_list, **context): def _jobs_are_old(cluster_name): """Return True if last slurm update in given cluster is older than 2 days.""" - jobs_are_old = False mongodb_filter = {"slurm.cluster_name": cluster_name} diff --git a/clockwork_web/templates/base.html b/clockwork_web/templates/base.html index bf79e7ea..8730dcd2 100644 --- a/clockwork_web/templates/base.html +++ b/clockwork_web/templates/base.html @@ -24,7 +24,7 @@ - + @@ -323,6 +323,12 @@

{% if previous_request_args['job_array'] is not none %} + {% endif %} + {% if previous_request_args['user_prop_name'] is not none %} + + {% endif %} + {% if previous_request_args['user_prop_content'] is not none %} + {% endif %}
@@ -334,6 +340,13 @@

{% endif %} + + {% if previous_request_args['user_prop_name'] is not none and previous_request_args['user_prop_content'] is not none %} + + User prop {{ previous_request_args['user_prop_name'] }}: "{{ previous_request_args['user_prop_content'] }}"     + + + {% endif %}

diff --git a/clockwork_web/templates/jobs_search.html b/clockwork_web/templates/jobs_search.html index 53077a6e..f225a7c0 100644 --- a/clockwork_web/templates/jobs_search.html +++ b/clockwork_web/templates/jobs_search.html @@ -101,6 +101,10 @@

JOBS

{% if (web_settings | check_web_settings_column_display(page_name, "job_array")) %} Job array {% endif %} + + {% if (web_settings | check_web_settings_column_display(page_name, "job_user_props")) %} + Job-user props + {% endif %} {% if (web_settings | check_web_settings_column_display(page_name, "job_name")) %} {% set sort_by = "name" %} @@ -193,6 +197,20 @@

JOBS

{% endif %} + + {% if (web_settings | check_web_settings_column_display(page_name, "job_user_props")) %} + + {% for D_user_prop_name, D_user_prop_content in D_job.get('job_user_props', {}).items() %} +

+ + {{ D_user_prop_name }}
+ {{ D_user_prop_content }} + +

+ {% endfor %} + + {% endif %} + {% if (web_settings | check_web_settings_column_display(page_name, "job_name")) %} {{D_job['slurm'].get("name", "")[0:20]}} diff --git a/clockwork_web/templates/settings.html b/clockwork_web/templates/settings.html index a2d28e3c..a9b04f51 100644 --- a/clockwork_web/templates/settings.html +++ b/clockwork_web/templates/settings.html @@ -279,6 +279,7 @@

{{ gettext("User settings %(mila_email_username)s", mila_email_username=curr {{ gettext("User (@mila.quebec)") }} {{ gettext("Job ID") }} {{ gettext("Job array") }} + {{ gettext("Job-user props") }} {{ gettext("Job name [:20]") }} {{ gettext("Job state") }} {{ gettext("Submit time") }} @@ -291,7 +292,7 @@

{{ gettext("User settings %(mila_email_username)s", mila_email_username=curr {% set page_name = "jobs_list" %} - {% for column_name in ["clusters", "user","job_id", "job_array", "job_name", "job_state", "submit_time", "start_time", "end_time", "links"] %} + {% for column_name in ["clusters", "user","job_id", "job_array", "job_user_props", "job_name", "job_state", "submit_time", "start_time", "end_time", "links"] %}
{% if (web_settings | check_web_settings_column_display(page_name, column_name)) %} diff --git a/clockwork_web/user.py b/clockwork_web/user.py index 27b3eff4..986de4c3 100644 --- a/clockwork_web/user.py +++ b/clockwork_web/user.py @@ -93,6 +93,10 @@ def boolean(value): self.web_settings.setdefault("column_display", {}).setdefault("jobs_list", {})[ "actions" ] = False + # By default, do not display column "job_user_props" in job tables. + self.web_settings.setdefault("column_display", {}).setdefault( + "jobs_list", {} + ).setdefault("job_user_props", False) def get_id(self): return self.mila_email_username diff --git a/scripts/store_huge_fake_data_in_db.py b/scripts/store_huge_fake_data_in_db.py new file mode 100644 index 00000000..b654eb2e --- /dev/null +++ b/scripts/store_huge_fake_data_in_db.py @@ -0,0 +1,271 @@ +""" +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_common/fake_data.json b/test_common/fake_data.json index b38bb1a2..29c18e58 100644 --- a/test_common/fake_data.json +++ b/test_common/fake_data.json @@ -5962,5 +5962,48 @@ "tensor_cores": 576, "tflops_fp32": 16.31 } + ], + "job_user_props": [ + { + "mila_email_username": "student00@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, + "cluster_name": "mila", + "props": { + "name": "je suis une user prop 2" + } + }, + { + "mila_email_username": "student00@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, + "cluster_name": "graham", + "props": { + "name": "je suis une user prop 3", + "name2": "je suis une user prop 4" + } + }, + { + "mila_email_username": "student00@mila.quebec", + "job_id": 613024, + "cluster_name": "graham", + "props": { + "name": "je suis une user prop 1" + } + } ] -} \ No newline at end of file +} diff --git a/test_common/fake_data.py b/test_common/fake_data.py index 5061ffc0..06f9cfdc 100644 --- a/test_common/fake_data.py +++ b/test_common/fake_data.py @@ -59,7 +59,7 @@ def populate_fake_data(db_insertion_point, json_file=None, mutate=False): # 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)], + [("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( @@ -70,8 +70,12 @@ def populate_fake_data(db_insertion_point, json_file=None, mutate=False): [("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", + ) - for k in ["users", "jobs", "nodes", "gpu"]: + for k in ["users", "jobs", "nodes", "gpu", "job_user_props"]: if k in E: for e in E[k]: db_insertion_point[k].insert_one(e) @@ -96,6 +100,11 @@ def cleanup_function(): for e in E["gpu"]: db_insertion_point["gpu"].delete_many({"name": e["name"]}) + for e in E["job_user_props"]: + copy_e = e.copy() + copy_e.pop("props") + db_insertion_point["job_user_props"].delete_many(copy_e) + for (k, sub, id_field) in [ ("jobs", "slurm", "job_id"), ("nodes", "slurm", "name"), diff --git a/test_common/jobs_test_helpers.py b/test_common/jobs_test_helpers.py index 5804a2ec..a1f81872 100644 --- a/test_common/jobs_test_helpers.py +++ b/test_common/jobs_test_helpers.py @@ -34,11 +34,8 @@ def helper_single_job_at_random(fake_data, cluster_name): def validator(D_job): for k1 in original_D_job: - assert k1 in ["slurm", "cw", "user"] - for k2 in original_D_job[k1]: - assert ( - D_job[k1][k2] == original_D_job[k1][k2] - ), f"{D_job}\n{original_D_job}" + assert k1 in ["slurm", "cw", "user", "job_user_props"] + assert D_job[k1] == original_D_job[k1], f"{D_job}\n{original_D_job}" return validator, job_id @@ -167,8 +164,7 @@ def validator(LD_jobs): # compare all the dicts one by one for (D_job, D_original_job) in zip(LD_jobs, LD_original_jobs): for k1 in D_original_job: - assert k1 in ["slurm", "cw", "user"] - for k2 in D_original_job[k1]: - assert D_job[k1][k2] == D_original_job[k1][k2] + assert k1 in ["slurm", "cw", "user", "job_user_props"] + assert D_job[k1] == D_original_job[k1] return validator