Merge pull request #9243 from OpenMined/eelco/configure-scenario-test…

…-size Configure scenario tests with env var
OpenMined · Sep 2, 2024 · 13ad1d0 · 13ad1d0
2 parents 4fd86b2 + c3d88d2
commit 13ad1d0
Show file tree

Hide file tree

Showing 8 changed files with 90 additions and 38 deletions.
diff --git a/notebooks/scenarios/bigquery/00-start-and-configure-server-and-admins.ipynb b/notebooks/scenarios/bigquery/00-start-and-configure-server-and-admins.ipynb
@@ -212,6 +212,13 @@
    "source": [
     "server.land()"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -230,7 +237,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.2"
+   "version": "3.10.13"
   }
  },
  "nbformat": 4,

diff --git a/notebooks/scenarios/bigquery/01-setup-datasite.ipynb b/notebooks/scenarios/bigquery/01-setup-datasite.ipynb
@@ -402,7 +402,9 @@
    "source": [
     "worker_pool_name = \"bigquery-pool\"\n",
     "custom_pool_pod_annotations = {\"bigquery-custom-pool\": \"Pod annotation for bigquery\"}\n",
-    "custom_pool_pod_labels = {\"bigquery-custom-pool\": \"Pod_label_for_bigquery\"}"
+    "custom_pool_pod_labels = {\"bigquery-custom-pool\": \"Pod_label_for_bigquery\"}\n",
+    "\n",
+    "num_workers = int(os.environ.get(\"NUM_TEST_WORKERS\", 1))"
    ]
   },
   {
@@ -428,7 +430,9 @@
    "outputs": [],
    "source": [
     "if environment == \"remote\":\n",
-    "    result = high_client.worker_pools.scale(number=2, pool_name=worker_pool_name)\n",
+    "    result = high_client.worker_pools.scale(\n",
+    "        number=num_workers, pool_name=worker_pool_name\n",
+    "    )\n",
     "    print(result)"
    ]
   },
@@ -476,6 +480,13 @@
    "source": [
     "server.land()"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -494,7 +505,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.2"
+   "version": "3.10.13"
   }
  },
  "nbformat": 4,

diff --git a/notebooks/scenarios/bigquery/011-users-emails-passwords.ipynb b/notebooks/scenarios/bigquery/011-users-emails-passwords.ipynb
@@ -129,14 +129,32 @@
     "# high_client.api.services.settings.disable_notifications()"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "9",
+   "metadata": {},
+   "source": [
+    "# Register users"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9",
+   "id": "10",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "num_users = int(os.environ.get(\"NUM_TEST_USERS\", 5))\n",
+    "print(f\"registering {num_users} users\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "11",
    "metadata": {},
    "outputs": [],
    "source": [
-    "num_users = 5\n",
     "users = []\n",
     "email_disable_index = 0\n",
     "reset_password_index = 1\n",
@@ -155,7 +173,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "10",
+   "id": "12",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -165,7 +183,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "11",
+   "id": "13",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -178,7 +196,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "12",
+   "id": "14",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -197,7 +215,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "13",
+   "id": "15",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -222,7 +240,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "14",
+   "id": "16",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -232,7 +250,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "15",
+   "id": "17",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -243,7 +261,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "16",
+   "id": "18",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -257,7 +275,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "17",
+   "id": "19",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -269,7 +287,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "18",
+   "id": "20",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -279,7 +297,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "19",
+   "id": "21",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -290,7 +308,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "20",
+   "id": "22",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -301,7 +319,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "21",
+   "id": "23",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -312,7 +330,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "22",
+   "id": "24",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -323,7 +341,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "23",
+   "id": "25",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -336,7 +354,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "24",
+   "id": "26",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -346,7 +364,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "25",
+   "id": "27",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -357,7 +375,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "26",
+   "id": "28",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -367,7 +385,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "27",
+   "id": "29",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -377,12 +395,20 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "28",
+   "id": "30",
    "metadata": {},
    "outputs": [],
    "source": [
     "server.land()"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "31",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -401,7 +427,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.2"
+   "version": "3.10.13"
   }
  },
  "nbformat": 4,

diff --git a/notebooks/scenarios/bigquery/02-configure-api.ipynb b/notebooks/scenarios/bigquery/02-configure-api.ipynb
@@ -535,7 +535,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.2"
+   "version": "3.10.13"
   }
  },
  "nbformat": 4,

diff --git a/notebooks/scenarios/bigquery/021-create-jobs.ipynb b/notebooks/scenarios/bigquery/021-create-jobs.ipynb
@@ -169,10 +169,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# maybe we should assign one of each first, and then randomly generate the rest\n",
-    "n_per_user = 2\n",
+    "num_jobs = int(os.environ.get(\"NUM_TEST_JOBS\", 10))\n",
     "\n",
-    "jobs = create_jobs(users, n_per_user=n_per_user)"
+    "jobs = create_jobs(users, total_jobs=num_jobs)"
    ]
   },
   {
@@ -182,7 +181,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "jobs"
+    "# stdlib\n",
+    "from collections import Counter\n",
+    "\n",
+    "counts = Counter([j.job_type for j in jobs])\n",
+    "for k, v in counts.most_common():\n",
+    "    print(f\"number of {k}: {v}\")"
    ]
   },
   {
@@ -220,7 +224,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "assert len(jobs) == len(users) * n_per_user\n",
+    "assert len(jobs) == num_jobs\n",
     "assert all(isinstance(j, TestJob) for j in jobs)\n",
     "assert all(job.client is not None for job in jobs)"
    ]
@@ -428,7 +432,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.2"
+   "version": "3.10.13"
   }
  },
  "nbformat": 4,

diff --git a/notebooks/scenarios/bigquery/04-do-review-requests.ipynb b/notebooks/scenarios/bigquery/04-do-review-requests.ipynb
@@ -253,7 +253,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.2"
+   "version": "3.10.13"
   }
  },
  "nbformat": 4,

diff --git a/notebooks/scenarios/bigquery/job_helpers.py b/notebooks/scenarios/bigquery/job_helpers.py
@@ -247,15 +247,14 @@ def create_job_many_columns(user: TestUser) -> TestJob:
     return job
 
 
-def create_job(user: TestUser) -> TestJob:
+def create_random_job(user: TestUser) -> TestJob:
     job_func = random.choice(create_job_functions)
     return job_func(user)
 
 
-def create_jobs(users: list[TestUser], n_per_user: int = 10) -> list[TestJob]:
+def create_jobs(users: list[TestUser], total_jobs: int = 10) -> list[TestJob]:
     jobs = []
     num_users = len(users)
-    total_jobs = n_per_user * num_users
     user_index = 0
     each_count = 0
     # keep making jobs until we have enough
@@ -266,7 +265,7 @@ def create_jobs(users: list[TestUser], n_per_user: int = 10) -> list[TestJob]:
             each_count += 1
         else:
             # otherwise lets get a random one
-            job_func = create_job
+            job_func = create_random_job
         # use the current index of user
         jobs.append(job_func(users[user_index]))