diff --git a/bigframes/core/global_session.py b/bigframes/core/global_session.py index 3187c5c11b..e70cdad59e 100644 --- a/bigframes/core/global_session.py +++ b/bigframes/core/global_session.py @@ -30,7 +30,7 @@ _global_session_state.thread_local_session = None -def _try_close_session(session): +def _try_close_session(session: bigframes.session.Session): """Try to close the session and warn if couldn't.""" try: session.close() diff --git a/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb b/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb index c0c3c58a3c..f56472f1b5 100644 --- a/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb +++ b/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb @@ -1217,6 +1217,15 @@ "Otherwise, you can uncomment the remaining cells and run them to delete the individual resources you created in this tutorial:" ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "bf.close_session()" + ] + }, { "cell_type": "code", "execution_count": 24, diff --git a/notebooks/getting_started/getting_started_bq_dataframes.ipynb b/notebooks/getting_started/getting_started_bq_dataframes.ipynb index b59ccbb8ac..e79db455fc 100644 --- a/notebooks/getting_started/getting_started_bq_dataframes.ipynb +++ b/notebooks/getting_started/getting_started_bq_dataframes.ipynb @@ -1704,6 +1704,16 @@ "Otherwise, you can uncomment the remaining cells and run them to delete the individual resources you created in this tutorial:" ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Delete the temporary cloud artifacts created during the bigframes session \n", + "bpd.close_session()" + ] + }, { "cell_type": "code", "execution_count": 25, diff --git a/notebooks/location/regionalized.ipynb b/notebooks/location/regionalized.ipynb index c383a22609..5a8239a42a 100644 --- a/notebooks/location/regionalized.ipynb +++ b/notebooks/location/regionalized.ipynb @@ -7,8 +7,9 @@ "source": [ "# README\n", "\n", - "This Notebook runs differently depending on the following environent variable:\n", - "1. BIGQUERY_LOCATION - can take values as per https://cloud.google.com/bigquery/docs/locations, e.g. `us`, `asia-east1`" + "This Notebook runs requiring the following environent variable:\n", + "1. GOOGLE_CLOUD_PROJECT - The google cloud project id.\n", + "1. BIGQUERY_LOCATION - can take values as per https://cloud.google.com/bigquery/docs/locations, e.g. `us`, `asia-east1`." ] }, { @@ -1420,8 +1421,8 @@ } ], "source": [ - "import bigframes.pandas as pd\n", - "help(pd.remote_function)" + "import bigframes.pandas as bpd\n", + "help(bpd.remote_function)" ] }, { @@ -1460,7 +1461,7 @@ } ], "source": [ - "@pd.remote_function([float], str, bigquery_connection='bigframes-rf-conn')\n", + "@bpd.remote_function([float], str, bigquery_connection='bigframes-rf-conn')\n", "def get_bucket(num):\n", " if not num: return \"NA\"\n", " boundary = 4000\n", @@ -2784,6 +2785,22 @@ "source": [ "model.to_gbq(f\"{DATASET}.penguins_model\", replace=True)" ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Clean Up" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "bpd.close_session()" + ] } ], "metadata": { diff --git a/notebooks/remote_functions/remote_function_usecases.ipynb b/notebooks/remote_functions/remote_function_usecases.ipynb index 9317e4b8fe..b897def4e8 100644 --- a/notebooks/remote_functions/remote_function_usecases.ipynb +++ b/notebooks/remote_functions/remote_function_usecases.ipynb @@ -25,7 +25,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# Setup" + "# Set Up" ] }, { @@ -1379,6 +1379,22 @@ "df1 = df.assign(duration_cat=df[\"duration_minutes\"].apply(duration_category))\n", "df1.peek()" ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Clean Up" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "bpd.close_session()" + ] } ], "metadata": { diff --git a/notebooks/remote_functions/remote_function_vertex_claude_model.ipynb b/notebooks/remote_functions/remote_function_vertex_claude_model.ipynb index 650bb92e50..a5769a2285 100644 --- a/notebooks/remote_functions/remote_function_vertex_claude_model.ipynb +++ b/notebooks/remote_functions/remote_function_vertex_claude_model.ipynb @@ -452,12 +452,21 @@ "df" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Clean Up" + ] + }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], - "source": [] + "source": [ + "bpd.close_session()" + ] } ], "metadata": { diff --git a/noxfile.py b/noxfile.py index 03c39f33c4..c868934d94 100644 --- a/noxfile.py +++ b/noxfile.py @@ -105,6 +105,7 @@ "system-3.9", "system-3.12", "cover", + "cleanup", ] # Error if a python version is missing @@ -944,25 +945,24 @@ def release_dry_run(session): def cleanup(session): """Clean up stale and/or temporary resources in the test project.""" google_cloud_project = os.getenv("GOOGLE_CLOUD_PROJECT") - if not google_cloud_project: - session.error( - "Set GOOGLE_CLOUD_PROJECT environment variable to run notebook session." - ) + cleanup_options = [] + if google_cloud_project: + cleanup_options.append(f"--project-id={google_cloud_project}") # Cleanup a few stale (more than 12 hours old) temporary cloud run # functions created by bigframems. This will help keeping the test GCP # project within the "Number of functions" quota # https://cloud.google.com/functions/quotas#resource_limits recency_cutoff_hours = 12 - cleanup_count_per_location = 10 + cleanup_count_per_location = 20 + cleanup_options.extend( + [ + f"--recency-cutoff={recency_cutoff_hours}", + "cleanup", + f"--number={cleanup_count_per_location}", + ] + ) session.install("-e", ".") - session.run( - "python", - "scripts/manage_cloud_functions.py", - f"--project-id={google_cloud_project}", - f"--recency-cutoff={recency_cutoff_hours}", - "cleanup", - f"--number={cleanup_count_per_location}", - ) + session.run("python", "scripts/manage_cloud_functions.py", *cleanup_options) diff --git a/scripts/manage_cloud_functions.py b/scripts/manage_cloud_functions.py index 33af8463c9..145e178f4d 100644 --- a/scripts/manage_cloud_functions.py +++ b/scripts/manage_cloud_functions.py @@ -153,6 +153,12 @@ def list_str(values): return [val for val in values.split(",") if val] +def get_project_from_environment(): + from google.cloud import bigquery + + return bigquery.Client().project + + if __name__ == "__main__": parser = argparse.ArgumentParser( description="Manage cloud functions created to serve bigframes remote functions." @@ -161,9 +167,10 @@ def list_str(values): "-p", "--project-id", type=str, - required=True, + required=False, action="store", - help="GCP project-id.", + help="GCP project-id. If not provided, the project-id resolved by the" + " BigQuery client from the user environment would be used.", ) parser.add_argument( "-r", @@ -212,4 +219,10 @@ def hours_to_timedelta(hrs): parser_cleanup.set_defaults(func=cleanup_gcfs) args = parser.parse_args(sys.argv[1:]) + if args.project_id is None: + args.project_id = get_project_from_environment() + if args.project_id is None: + raise ValueError( + "Could not resolve a project. Plese set it via --project-id option." + ) args.func(args)