diff --git a/notebooks/README.md b/notebooks/README.md index d00c7fcd337..845a421ba58 100644 --- a/notebooks/README.md +++ b/notebooks/README.md @@ -1,82 +1,12 @@ -# AI Platform Notebooks +# Google Cloud Notebooks & Workbench Samples -[AI Platform Notebooks](https://cloud.google.com/ai-platform-notebooks) -is a managed service that offers an integrated and secure JupyterLab -environment for data scientists and machine learning developers to -experiment, develop, and deploy models into production. +This directory contains samples and tutorials for managing **Workbench instances** and executing Jupyter notebooks on Google Cloud. -This directory contains code samples and -tutorials. The tutorials are divided into 2 sections: -- [API code samples](samples) -- [tutorials](tutorials) +Workbench instances provide enterprises with a managed, secure JupyterLab environment that protects organizational data while enabling centralized administration and configuration. Data scientists can experiment, develop, and deploy machine learning and analytics workloads on Google Cloud while maintaining their organization's required security posture. -## Notebooks API +## Directory Structure -In [this](samples) folder you will find Python code samples to interact -with the -[API](https://cloud.google.com/ai-platform/notebooks/docs/reference/rest) - -## Tutorials - -**Run in Notebooks** - -1. Create a Notebook - [instance](https://cloud.google.com/ai-platform/notebooks/docs#how-to) -2. Clone this repository via Jupyter console or using the Git clone - button - -```bash -git clone https://github.com/GoogleCloudPlatform/python-docs-samples.git -``` - -3. Open the tutorials folder and open the Jupyter notebooks - - -**Run locally** - -1. Install Jupyter notebooks - ([instructions](https://jupyter.org/install)) -1. Install the dependencies in the [requirements.txt](./requirements.txt) file ([instructions below](#install-the-dependencies)) -1. Registered the `google-cloud-bigquery` magic commands ([instructions below](#register-magics-and-configure-matplotlib)) -1. Set `matplotlib` to render inline ([instructions below](#register-magics-and-configure-matplotlib)) - -### Install the dependencies - -Install the dependencies with the following command: - - pip install --upgrade -r requirements.txt - -### Register magics and configure matplotlib - -You can either perform these set up steps in a single notebook, or add the -steps to your IPython configuration file to apply to all notebooks. - -### Perform set up steps within a notebook - -To perform the set up steps for a single notebook, run the following commands -in your notebook to register the BigQuery magic commands and set `matplotlib` -to render inline: -```python -%load_ext google.cloud.bigquery -%matplotlib inline -``` - -### Perform set up steps in your IPython configuration file - -To perform the set up steps implicitly for all of your notebooks, add the -following code to your `ipython_config.py` file to register the BigQuery magic -commands and set `matplotlib` to render inline: -```python -c = get_config() - -# Register magic commands -c.InteractiveShellApp.extensions = [ - 'google.cloud.bigquery', -] - -# Enable matplotlib renderings to render inline in the notebook. -c.InteractiveShellApp.matplotlib = 'inline' -``` -See -[IPython documentation](https://ipython.readthedocs.io/en/stable/config/intro.html) -for more information about IPython configuration. +- **[`v2/`](v2/)**: Atomic Python client library samples (`google-cloud-notebooks` / `google.cloud.notebooks_v2`) for creating, configuring, hardening, and managing **Workbench instances**, including enterprise security controls (private VPC networking, single-user least-privilege service accounts, CMEK encryption, Shielded VM configuration, and IAM access policies). +- **[`rendered/`](rendered/)**: Rendered HTML views of tutorial notebooks. +- **[`tutorials/bigquery/`](tutorials/bigquery/)**: Jupyter notebook tutorials demonstrating BigQuery integration from a notebook environment. +- **[`tutorials/vision/`](tutorials/vision/)**: Jupyter notebook tutorials demonstrating Cloud Vision API integration from a notebook environment. diff --git a/notebooks/samples/README.md b/notebooks/samples/README.md deleted file mode 100644 index 7a92a46d3d0..00000000000 --- a/notebooks/samples/README.md +++ /dev/null @@ -1,39 +0,0 @@ -# API code samples - -Sample command-line programs for interacting with the Notebooks API. - -## Prerequisites to run locally: - -* [Python 3](https://www.python.org/downloads/) -* [pip](https://pypi.python.org/pypi/pip) - -Go to the [Google Cloud Console](https://console.cloud.google.com). - -Under API Manager, search for the Google Cloud Notebooks API and enable -it. You can also enable API via Command Line: - -```bash -gcloud services enable notebooks.googleapis.com -``` - -## Set Up Your Local Dev Environment - -To install, run the following commands. If you want to use -[virtualenv](https://virtualenv.readthedocs.org/en/latest/) -(recommended), run the commands within a virtualenv. - -```bash -pip install -r requirements.txt -``` - -**Note:** While this sample demonstrates interacting with Notebooks via -the API, the functionality demonstrated here could also be accomplished -using the Cloud Console or the -[gcloud CLI](https://cloud.google.com/sdk/gcloud/reference/beta/notebooks) - -## Authentication - -Please see the -[Google cloud authentication guide](https://cloud.google.com/docs/authentication/). -The recommended approach to running these samples is a Service Account -with a JSON key. diff --git a/notebooks/samples/requirements-test.txt b/notebooks/samples/requirements-test.txt deleted file mode 100644 index c9e154ba440..00000000000 --- a/notebooks/samples/requirements-test.txt +++ /dev/null @@ -1 +0,0 @@ -pytest==9.0.3; python_version >= "3.10" diff --git a/notebooks/samples/requirements.txt b/notebooks/samples/requirements.txt deleted file mode 100644 index 76dc107a22b..00000000000 --- a/notebooks/samples/requirements.txt +++ /dev/null @@ -1 +0,0 @@ -google-api-python-client==2.131.0 diff --git a/notebooks/tutorials/cloud-ml-engine/Training and prediction with scikit-learn.ipynb b/notebooks/tutorials/cloud-ml-engine/Training and prediction with scikit-learn.ipynb deleted file mode 100644 index d805d0cb1fe..00000000000 --- a/notebooks/tutorials/cloud-ml-engine/Training and prediction with scikit-learn.ipynb +++ /dev/null @@ -1,570 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Training and prediction with scikit-learn\n", - "\n", - "This notebook demonstrates how to use AI Platform to train a simple classification model using `scikit-learn`, and then deploy the model to get predictions.\n", - "\n", - "You train the model to predict a person's income level based on the [Census Income data set](https://archive.ics.uci.edu/ml/datasets/Census+Income).\n", - "\n", - "Before you jump in, let’s cover some of the different tools you’ll be using:\n", - "\n", - "+ [AI Platform](https://cloud.google.com/ml-engine/) is a managed service that enables you to easily build machine learning models that work on any type of data, of any size.\n", - "\n", - "+ [Cloud Storage](https://cloud.google.com/storage/) is a unified object storage for developers and enterprises, from live data serving to data analytics/ML to data archiving.\n", - "\n", - "+ [Cloud SDK](https://cloud.google.com/sdk/) is a command line tool which allows you to interact with Google Cloud products. This notebook introduces several `gcloud` and `gsutil` commands, which are part of the Cloud SDK. Note that shell commands in a notebook must be prepended with a `!`." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Set up your environment\n", - "\n", - "### Enable the required APIs\n", - "\n", - "In order to use AI Platform, confirm that the required APIs are enabled:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "!gcloud services enable ml.googleapis.com\n", - "!gcloud services enable compute.googleapis.com" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Create a storage bucket\n", - "Buckets are the basic containers that hold your data. Everything that you store in Cloud Storage must be contained in a bucket. You can use buckets to organize your data and control access to your data.\n", - "\n", - "Start by defining a globally unique name.\n", - "\n", - "For more information about naming buckets, see [Bucket name requirements](https://cloud.google.com/storage/docs/naming#requirements)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "BUCKET_NAME = \"your-new-bucket\"" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the examples below, the `BUCKET_NAME` variable is referenced in the commands using `$`.\n", - "\n", - "Create the new bucket with the `gsutil mb` command:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "!gcloud storage buckets create gs://$BUCKET_NAME" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### About the data\n", - "\n", - "The [Census Income Data Set](https://archive.ics.uci.edu/ml/datasets/Census+Income) that this sample\n", - "uses for training is provided by the [UC Irvine Machine Learning\n", - "Repository](https://archive.ics.uci.edu/ml/datasets/).\n", - "\n", - "Census data courtesy of: Lichman, M. (2013). UCI Machine Learning Repository http://archive.ics.uci.edu/ml. Irvine, CA: University of California, School of Information and Computer Science. This dataset is publicly available for anyone to use under the following terms provided by the Dataset Source - http://archive.ics.uci.edu/ml - and is provided \"AS IS\" without any warranty, express or implied, from Google. Google disclaims all liability for any damages, direct or indirect, resulting from the use of the dataset.\n", - "\n", - "The data used in this tutorial is located in a public Cloud Storage bucket:\n", - "\n", - " gs://cloud-samples-data/ml-engine/sklearn/census_data/ \n", - "\n", - "The training file is `adult.data` ([download](https://storage.googleapis.com/cloud-samples-data/ml-engine/sklearn/census_data/adult.data)) and the evaluation file is `adult.test` ([download](https://storage.googleapis.com/cloud-samples-data/ml-engine/sklearn/census_data/adult.test)). The evaluation file is not used in this tutorial." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Create training application package\n", - "\n", - "The easiest (and recommended) way to create a training application package is to use `gcloud` to package and upload the application when you submit your training job. This method allows you to create a very simple file structure with only two files. For this tutorial, the file structure of your training application package should appear similar to the following:\n", - "\n", - " census_training/\n", - " __init__.py\n", - " train.py" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create a directory locally:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "!mkdir census_training" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create a blank file named `__init__.py`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "!touch ./census_training/__init__.py" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Save training code in one Python file in the `census_training` directory. The following cell writes a training file to the `census_training` directory. The training file performs the following operations:\n", - "+ Loads the data into a pandas `DataFrame` that can be used by `scikit-learn`\n", - "+ Fits the model is against the training data\n", - "+ Exports the model with the [Python `pickle` library](https://docs.python.org/3/library/pickle.html)\n", - "\n", - "The following model training code is not executed within this notebook. Instead, it is saved to a Python file and packaged as a Python module that runs on AI Platform after you submit the training job." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "%%writefile ./census_training/train.py\n", - "import argparse\n", - "import pickle\n", - "import pandas as pd\n", - "\n", - "from google.cloud import storage\n", - "\n", - "from sklearn.ensemble import RandomForestClassifier\n", - "from sklearn.feature_selection import SelectKBest\n", - "from sklearn.pipeline import FeatureUnion\n", - "from sklearn.pipeline import Pipeline\n", - "from sklearn.preprocessing import LabelBinarizer\n", - "\n", - "parser = argparse.ArgumentParser()\n", - "parser.add_argument(\"--bucket-name\", help=\"The bucket name\", required=True)\n", - "\n", - "arguments, unknown = parser.parse_known_args()\n", - "bucket_name = arguments.bucket_name\n", - "\n", - "# Define the format of your input data, including unused columns.\n", - "# These are the columns from the census data files.\n", - "COLUMNS = (\n", - " 'age',\n", - " 'workclass',\n", - " 'fnlwgt',\n", - " 'education',\n", - " 'education-num',\n", - " 'marital-status',\n", - " 'occupation',\n", - " 'relationship',\n", - " 'race',\n", - " 'sex',\n", - " 'capital-gain',\n", - " 'capital-loss',\n", - " 'hours-per-week',\n", - " 'native-country',\n", - " 'income-level'\n", - ")\n", - "\n", - "# Categorical columns are columns that need to be turned into a numerical value\n", - "# to be used by scikit-learn\n", - "CATEGORICAL_COLUMNS = (\n", - " 'workclass',\n", - " 'education',\n", - " 'marital-status',\n", - " 'occupation',\n", - " 'relationship',\n", - " 'race',\n", - " 'sex',\n", - " 'native-country'\n", - ")\n", - "\n", - "# Create a Cloud Storage client to download the census data\n", - "storage_client = storage.Client()\n", - "\n", - "# Download the data\n", - "public_bucket = storage_client.bucket('cloud-samples-data')\n", - "blob = public_bucket.blob('ml-engine/sklearn/census_data/adult.data')\n", - "blob.download_to_filename('adult.data')\n", - "\n", - "# Load the training census dataset\n", - "with open(\"./adult.data\", \"r\") as train_data:\n", - " raw_training_data = pd.read_csv(train_data, header=None, names=COLUMNS)\n", - " # Removing the whitespaces in categorical features\n", - " for col in CATEGORICAL_COLUMNS:\n", - " raw_training_data[col] = raw_training_data[col].apply(lambda x: str(x).strip())\n", - "\n", - "# Remove the column we are trying to predict ('income-level') from our features\n", - "# list and convert the DataFrame to a lists of lists\n", - "train_features = raw_training_data.drop(\"income-level\", axis=1).values.tolist()\n", - "# Create our training labels list, convert the DataFrame to a lists of lists\n", - "train_labels = (raw_training_data[\"income-level\"] == \" >50K\").values.tolist()\n", - "\n", - "# Since the census data set has categorical features, we need to convert\n", - "# them to numerical values. We'll use a list of pipelines to convert each\n", - "# categorical column and then use FeatureUnion to combine them before calling\n", - "# the RandomForestClassifier.\n", - "categorical_pipelines = []\n", - "\n", - "# Each categorical column needs to be extracted individually and converted to a\n", - "# numerical value. To do this, each categorical column will use a pipeline that\n", - "# extracts one feature column via SelectKBest(k=1) and a LabelBinarizer() to\n", - "# convert the categorical value to a numerical one. A scores array (created\n", - "# below) will select and extract the feature column. The scores array is\n", - "# created by iterating over the columns and checking if it is a\n", - "# categorical column.\n", - "for i, col in enumerate(COLUMNS[:-1]):\n", - " if col in CATEGORICAL_COLUMNS:\n", - " # Create a scores array to get the individual categorical column.\n", - " # Example:\n", - " # data = [\n", - " # 39, 'State-gov', 77516, 'Bachelors', 13, 'Never-married',\n", - " # 'Adm-clerical', 'Not-in-family', 'White', 'Male', 2174, 0,\n", - " # 40, 'United-States'\n", - " # ]\n", - " # scores = [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n", - " #\n", - " # Returns: [['State-gov']]\n", - " # Build the scores array\n", - " scores = [0] * len(COLUMNS[:-1])\n", - " # This column is the categorical column we want to extract.\n", - " scores[i] = 1\n", - " skb = SelectKBest(k=1)\n", - " skb.scores_ = scores\n", - " # Convert the categorical column to a numerical value\n", - " lbn = LabelBinarizer()\n", - " r = skb.transform(train_features)\n", - " lbn.fit(r)\n", - " # Create the pipeline to extract the categorical feature\n", - " categorical_pipelines.append(\n", - " (\n", - " 'categorical-{}'.format(i), \n", - " Pipeline([\n", - " ('SKB-{}'.format(i), skb),\n", - " ('LBN-{}'.format(i), lbn)])\n", - " )\n", - " )\n", - "\n", - "# Create pipeline to extract the numerical features\n", - "skb = SelectKBest(k=6)\n", - "# From COLUMNS use the features that are numerical\n", - "skb.scores_ = [1, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0]\n", - "categorical_pipelines.append((\"numerical\", skb))\n", - "\n", - "# Combine all the features using FeatureUnion\n", - "preprocess = FeatureUnion(categorical_pipelines)\n", - "\n", - "# Create the classifier\n", - "classifier = RandomForestClassifier()\n", - "\n", - "# Transform the features and fit them to the classifier\n", - "classifier.fit(preprocess.transform(train_features), train_labels)\n", - "\n", - "# Create the overall model as a single pipeline\n", - "pipeline = Pipeline([(\"union\", preprocess), (\"classifier\", classifier)])\n", - "\n", - "# Create the model file\n", - "# It is required to name the model file \"model.pkl\" if you are using pickle\n", - "model_filename = \"model.pkl\"\n", - "with open(model_filename, \"wb\") as model_file:\n", - " pickle.dump(pipeline, model_file)\n", - "\n", - "# Upload the model to Cloud Storage\n", - "bucket = storage_client.bucket(bucket_name)\n", - "blob = bucket.blob(model_filename)\n", - "blob.upload_from_filename(model_filename)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Submit the training job\n", - "\n", - "In this section, you use [`gcloud ai-platform jobs submit training`](https://cloud.google.com/sdk/gcloud/reference/ai-platform/jobs/submit/training) to submit your training job. The `--` argument passed to the command is a separator; anything after the separator will be passed to the Python code as input arguments.\n", - "\n", - "For more information about the arguments preceeding the separator, run the following:\n", - "\n", - " gcloud ai-platform jobs submit training --help\n", - "\n", - "The argument given to the python script is `--bucket-name`. The `--bucket-name` argument is used to specify the name of the bucket to save the model file." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "scrolled": true - }, - "outputs": [], - "source": [ - "import time\n", - "\n", - "# Define a timestamped job name\n", - "JOB_NAME = \"census_training_{}\".format(int(time.time()))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "scrolled": true - }, - "outputs": [], - "source": [ - "# Submit the training job:\n", - "!gcloud ai-platform jobs submit training $JOB_NAME \\\n", - " --job-dir gs://$BUCKET_NAME/census_job_dir \\\n", - " --package-path ./census_training \\\n", - " --module-name census_training.train \\\n", - " --region us-central1 \\\n", - " --runtime-version=1.12 \\\n", - " --python-version=3.5 \\\n", - " --scale-tier BASIC \\\n", - " --stream-logs \\\n", - " -- \\\n", - " --bucket-name $BUCKET_NAME" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Verify model file in Cloud Storage\n", - "\n", - "View the contents of the destination model directory to verify that your model file has been uploaded to Cloud Storage.\n", - "\n", - "Note: The model can take a few minutes to train and show up in Cloud Storage." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "!gcloud storage ls gs://$BUCKET_NAME/" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Serve the model" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Once the model is successfully created and trained, you can serve it. A model can have different versions. In order to serve the model, create a model and version in AI Platform.\n", - "\n", - "Define the model and version names:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "MODEL_NAME = \"CensusPredictor\"\n", - "VERSION_NAME = \"census_predictor_{}\".format(int(time.time()))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create the model in AI Platform:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "!gcloud ai-platform models create $MODEL_NAME --regions us-central1" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create a version that points to your model file in Cloud Storage:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "!gcloud ai-platform versions create $VERSION_NAME \\\n", - " --model=$MODEL_NAME \\\n", - " --framework=scikit-learn \\\n", - " --origin=gs://$BUCKET_NAME/ \\\n", - " --python-version=3.5 \\\n", - " --runtime-version=1.12" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Make predictions\n", - "\n", - "### Format data for prediction\n", - "\n", - "Before you send an online prediction request, you must format your test data to prepare it for use by the AI Platform prediction service. Make sure that the format of your input instances matches what your model expects.\n", - "\n", - "Create an `input.json` file with each input instance on a separate line. The following example uses ten data instances. Note that the format of input instances needs to match what your model expects. In this example, the Census model requires 14 features, so your input must be a matrix of shape (`num_instances, 14`)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Define a name for the input file\n", - "INPUT_FILE = \"./census_training/input.json\"" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "scrolled": true - }, - "outputs": [], - "source": [ - "%%writefile $INPUT_FILE\n", - "[25, \"Private\", 226802, \"11th\", 7, \"Never-married\", \"Machine-op-inspct\", \"Own-child\", \"Black\", \"Male\", 0, 0, 40, \"United-States\"]\n", - "[38, \"Private\", 89814, \"HS-grad\", 9, \"Married-civ-spouse\", \"Farming-fishing\", \"Husband\", \"White\", \"Male\", 0, 0, 50, \"United-States\"]\n", - "[28, \"Local-gov\", 336951, \"Assoc-acdm\", 12, \"Married-civ-spouse\", \"Protective-serv\", \"Husband\", \"White\", \"Male\", 0, 0, 40, \"United-States\"]\n", - "[44, \"Private\", 160323, \"Some-college\", 10, \"Married-civ-spouse\", \"Machine-op-inspct\", \"Husband\", \"Black\", \"Male\", 7688, 0, 40, \"United-States\"]\n", - "[18, \"?\", 103497, \"Some-college\", 10, \"Never-married\", \"?\", \"Own-child\", \"White\", \"Female\", 0, 0, 30, \"United-States\"]\n", - "[34, \"Private\", 198693, \"10th\", 6, \"Never-married\", \"Other-service\", \"Not-in-family\", \"White\", \"Male\", 0, 0, 30, \"United-States\"]\n", - "[29, \"?\", 227026, \"HS-grad\", 9, \"Never-married\", \"?\", \"Unmarried\", \"Black\", \"Male\", 0, 0, 40, \"United-States\"]\n", - "[63, \"Self-emp-not-inc\", 104626, \"Prof-school\", 15, \"Married-civ-spouse\", \"Prof-specialty\", \"Husband\", \"White\", \"Male\", 3103, 0, 32, \"United-States\"]\n", - "[24, \"Private\", 369667, \"Some-college\", 10, \"Never-married\", \"Other-service\", \"Unmarried\", \"White\", \"Female\", 0, 0, 40, \"United-States\"]\n", - "[55, \"Private\", 104996, \"7th-8th\", 4, \"Married-civ-spouse\", \"Craft-repair\", \"Husband\", \"White\", \"Male\", 0, 0, 10, \"United-States\"]" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Send the online prediction request\n", - "\n", - "The prediction results return `True` if the person's income is predicted to be greater than $50,000 per year, and `False` otherwise. The output of the command below may appear similar to the following:\n", - "\n", - " [False, False, False, True, False, False, False, False, False, False]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "scrolled": true - }, - "outputs": [], - "source": [ - "!gcloud ai-platform predict --model $MODEL_NAME --version \\\n", - " $VERSION_NAME --json-instances $INPUT_FILE" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Clean up" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To delete all resources you created in this tutorial, run the following commands:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "scrolled": true - }, - "outputs": [], - "source": [ - "# Delete the model version\n", - "!gcloud ai-platform versions delete $VERSION_NAME --model=$MODEL_NAME --quiet\n", - "\n", - "# Delete the model\n", - "!gcloud ai-platform models delete $MODEL_NAME --quiet\n", - "\n", - "# Delete the bucket and contents\n", - "!gcloud storage rm --recursive gs://$BUCKET_NAME\n", - "\n", - "# Delete the local files created by the tutorial\n", - "!rm -rf census_training" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.4" - } - }, - "nbformat": 4, - "nbformat_minor": 2 -} diff --git a/notebooks/v2/README.md b/notebooks/v2/README.md new file mode 100644 index 00000000000..7b8631fb033 --- /dev/null +++ b/notebooks/v2/README.md @@ -0,0 +1,60 @@ +# Workbench Instance Samples (Notebooks API v2) + +This directory contains atomic Python code samples demonstrating how to manage **Workbench instances** using the `google-cloud-notebooks` (`google.cloud.notebooks_v2`) client library. + +Workbench instances target enterprises that want to manage a secure JupyterLab environment that keeps their data safe while allowing centralized administration and configuration. Data scientists can experiment, develop, and deploy workloads on Google Cloud while maintaining their organization's desired security posture. + +## Prerequisites + +1. Enable the **Notebooks API** (`notebooks.googleapis.com`) in your Google Cloud project: + ```bash + gcloud services enable notebooks.googleapis.com + ``` +2. Configure Application Default Credentials (ADC): + ```bash + gcloud auth application-default login + ``` +3. Create a virtual environment and install dependencies: + ```bash + python3 -m venv .venv + .venv/bin/pip install -r requirements.txt + ``` + +## Samples + +### Core Instance Lifecycle & Read Operations +- [`list_instances.py`](list_instances.py) — Lists Workbench instances in a zone (`notebooks_v2_list_instances`). +- [`get_instance.py`](get_instance.py) — Gets metadata and health state for a Workbench instance (`notebooks_v2_get_instance`). +- [`create_instance.py`](create_instance.py) — Creates a new Workbench instance with default VM settings (`notebooks_v2_create_instance`). +- [`stop_instance.py`](stop_instance.py) — Stops a running Workbench instance (`notebooks_v2_stop_instance`). +- [`start_instance.py`](start_instance.py) — Starts a stopped Workbench instance (`notebooks_v2_start_instance`). +- [`reset_instance.py`](reset_instance.py) — Resets (reboots) a Workbench instance (`notebooks_v2_reset_instance`). +- [`delete_instance.py`](delete_instance.py) — Deletes a Workbench instance (`notebooks_v2_delete_instance`). + +### Centralized Administration, Maintenance & Diagnostics +- [`update_instance_machine_type.py`](update_instance_machine_type.py) — Updates the Compute Engine machine type of a stopped Workbench instance (`notebooks_v2_update_instance_machine_type`). +- [`update_instance_metadata.py`](update_instance_metadata.py) — Updates custom metadata keys on a Workbench instance (`notebooks_v2_update_instance_metadata`). +- [`update_instance_shielded_config.py`](update_instance_shielded_config.py) — Updates the Shielded VM configuration on a stopped Workbench instance (`notebooks_v2_update_instance_shielded_config`). +- [`check_instance_upgradability.py`](check_instance_upgradability.py) — Checks whether a Workbench instance can be upgraded (`notebooks_v2_check_instance_upgradability`). +- [`upgrade_instance.py`](upgrade_instance.py) — Upgrades a Workbench instance to the latest environment version (`notebooks_v2_upgrade_instance`). +- [`rollback_instance.py`](rollback_instance.py) — Rolls back a Workbench instance to a previous snapshot and revision (`notebooks_v2_rollback_instance`). +- [`diagnose_instance.py`](diagnose_instance.py) — Runs a diagnostic bundle export to Cloud Storage for a Workbench instance (`notebooks_v2_diagnose_instance`). + +### Enterprise Security Posture, Specialized Compute & Access Control +- [`create_instance_gpu.py`](create_instance_gpu.py) — Creates a GPU-accelerated Workbench instance (`notebooks_v2_create_instance_gpu`). +- [`create_instance_custom_container.py`](create_instance_custom_container.py) — Creates a Workbench instance from a custom container image (`notebooks_v2_create_instance_custom_container`). +- [`create_instance_cmek.py`](create_instance_cmek.py) — Creates a Workbench instance encrypted with a customer-managed encryption key (`notebooks_v2_create_instance_cmek`). +- [`create_instance_confidential.py`](create_instance_confidential.py) — Creates a Shielded VM Workbench instance (`notebooks_v2_create_instance_confidential`). +- [`create_instance_private_vpc.py`](create_instance_private_vpc.py) — Creates a private-IP Workbench instance in a custom VPC network and subnet (`notebooks_v2_create_instance_private_vpc`). +- [`create_instance_single_user.py`](create_instance_single_user.py) — Creates a single-user Workbench instance with a custom least-privilege service account (`notebooks_v2_create_instance_single_user`). +- [`create_instance_hardened.py`](create_instance_hardened.py) — Creates a security-hardened Workbench instance disabling root, downloads, terminal, and public IP (`notebooks_v2_create_instance_hardened`). +- [`get_instance_iam_policy.py`](get_instance_iam_policy.py) — Gets the IAM policy for a Workbench instance (`notebooks_v2_get_instance_iam_policy`). +- [`set_instance_iam_policy.py`](set_instance_iam_policy.py) — Sets an IAM policy binding on a Workbench instance using read-modify-write (`notebooks_v2_set_instance_iam_policy`). +- [`check_instance_iam_permissions.py`](check_instance_iam_permissions.py) — Checks IAM permissions on a Workbench instance (`notebooks_v2_check_instance_iam_permissions`). + +## Running Tests + +```bash +.venv/bin/pip install -r requirements-test.txt +.venv/bin/pytest +``` diff --git a/notebooks/v2/__init__.py b/notebooks/v2/__init__.py new file mode 100644 index 00000000000..58d482ea386 --- /dev/null +++ b/notebooks/v2/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/notebooks/v2/check_instance_iam_permissions.py b/notebooks/v2/check_instance_iam_permissions.py new file mode 100644 index 00000000000..d43de2d2253 --- /dev/null +++ b/notebooks/v2/check_instance_iam_permissions.py @@ -0,0 +1,54 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_check_instance_iam_permissions] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +from google.iam.v1 import iam_policy_pb2 + + +def check_instance_iam_permissions( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Checks IAM permissions on a Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + request = iam_policy_pb2.TestIamPermissionsRequest( + resource=name, + permissions=["notebooks.instances.get"], + ) + + try: + response = client.test_iam_permissions(request=request) + print(f"Resource ID: {name}") + print(f"Allowed Permissions: {list(response.permissions)}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_check_instance_iam_permissions] diff --git a/notebooks/v2/check_instance_upgradability.py b/notebooks/v2/check_instance_upgradability.py new file mode 100644 index 00000000000..c544c8c0df1 --- /dev/null +++ b/notebooks/v2/check_instance_upgradability.py @@ -0,0 +1,55 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_check_instance_upgradability] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def check_instance_upgradability( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Checks whether a Workbench instance can be upgraded. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to check. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + request = notebooks_v2.CheckInstanceUpgradabilityRequest( + notebook_instance=name, + ) + + try: + response = client.check_instance_upgradability(request=request) + print(f"Resource ID: {name}") + print(f"Upgrade Image: {response.upgrade_image}") + print(f"Upgrade Info: {response.upgrade_info}") + print(f"Upgrade Version: {response.upgrade_version}") + print(f"Upgradeable: {response.upgradeable}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_check_instance_upgradability] diff --git a/notebooks/v2/create_instance.py b/notebooks/v2/create_instance.py new file mode 100644 index 00000000000..96c7dd77ddb --- /dev/null +++ b/notebooks/v2/create_instance.py @@ -0,0 +1,67 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_create_instance] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def create_instance( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Creates a Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone for the new instance. + instance_id: The ID to assign to the new Workbench instance. + """ + client = notebooks_v2.NotebookServiceClient() + parent = f"projects/{project_id}/locations/{location_id}" + gce_setup = notebooks_v2.GceSetup(machine_type="e2-standard-4") + instance = notebooks_v2.Instance(gce_setup=gce_setup) + request = notebooks_v2.CreateInstanceRequest( + parent=parent, + instance_id=instance_id, + instance=instance, + ) + + try: + operation = client.create_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + # The operation completes once the instance is accepted, before the + # VM finishes provisioning, so the state is usually PROVISIONING here. + response = operation.result(timeout=300) + print(f"Resource ID: {response.name}") + print(f"Creation Date: {response.create_time}") + print(f"Machine Type: {response.gce_setup.machine_type}") + print(f"Proxy URI: {response.proxy_uri}") + print(f"State: {response.state.name}") + print( + "The instance is still starting up; call get_instance until the " + "state is ACTIVE before connecting to JupyterLab." + ) + except exceptions.AlreadyExists as err: + print( + f"error: Instance '{instance_id}' already exists in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_create_instance] diff --git a/notebooks/v2/create_instance_cmek.py b/notebooks/v2/create_instance_cmek.py new file mode 100644 index 00000000000..ae34ea45b64 --- /dev/null +++ b/notebooks/v2/create_instance_cmek.py @@ -0,0 +1,93 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_create_instance_cmek] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def create_instance_cmek( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-cmek-instance", + key_ring_id: str = "example-key-ring", + crypto_key_id: str = "example-key", +) -> None: + """Creates a Workbench instance encrypted with a customer-managed key. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone for the new instance. + instance_id: The ID to assign to the new Workbench instance. + key_ring_id: The Cloud KMS key ring ID in the instance region. + crypto_key_id: The Cloud KMS crypto key ID used for disk encryption. + """ + client = notebooks_v2.NotebookServiceClient() + parent = f"projects/{project_id}/locations/{location_id}" + parts = location_id.split("-") + region_id = "-".join(parts[:-1]) if len(parts) > 2 else location_id + kms_key = ( + f"projects/{project_id}/locations/{region_id}/" + f"keyRings/{key_ring_id}/cryptoKeys/{crypto_key_id}" + ) + boot_disk = notebooks_v2.BootDisk( + disk_size_gb=150, + disk_type=notebooks_v2.DiskType.PD_SSD, + disk_encryption=notebooks_v2.DiskEncryption.CMEK, + kms_key=kms_key, + ) + data_disk = notebooks_v2.DataDisk( + disk_size_gb=100, + disk_type=notebooks_v2.DiskType.PD_SSD, + disk_encryption=notebooks_v2.DiskEncryption.CMEK, + kms_key=kms_key, + ) + gce_setup = notebooks_v2.GceSetup( + machine_type="e2-standard-4", + boot_disk=boot_disk, + data_disks=[data_disk], + ) + instance = notebooks_v2.Instance(gce_setup=gce_setup) + request = notebooks_v2.CreateInstanceRequest( + parent=parent, + instance_id=instance_id, + instance=instance, + ) + + try: + operation = client.create_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + # The operation completes once the instance is accepted, before the + # VM finishes provisioning, so the state is usually PROVISIONING here. + response = operation.result(timeout=300) + boot_encryption = response.gce_setup.boot_disk.disk_encryption.name + print(f"Resource ID: {response.name}") + print(f"Creation Date: {response.create_time}") + print(f"Boot Disk Encryption: {boot_encryption}") + print(f"State: {response.state.name}") + print( + "The instance is still starting up; call get_instance until the " + "state is ACTIVE before connecting to JupyterLab." + ) + except exceptions.AlreadyExists as err: + print( + f"error: Instance '{instance_id}' already exists in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_create_instance_cmek] diff --git a/notebooks/v2/create_instance_confidential.py b/notebooks/v2/create_instance_confidential.py new file mode 100644 index 00000000000..5a4ba70eac5 --- /dev/null +++ b/notebooks/v2/create_instance_confidential.py @@ -0,0 +1,75 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_create_instance_confidential] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def create_instance_confidential( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-confidential-instance", +) -> None: + """Creates a Shielded VM Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone for the new instance. + instance_id: The ID to assign to the new Workbench instance. + """ + client = notebooks_v2.NotebookServiceClient() + parent = f"projects/{project_id}/locations/{location_id}" + shielded_config = notebooks_v2.ShieldedInstanceConfig( + enable_secure_boot=True, + enable_vtpm=True, + enable_integrity_monitoring=True, + ) + gce_setup = notebooks_v2.GceSetup( + machine_type="n2d-standard-4", + shielded_instance_config=shielded_config, + ) + instance = notebooks_v2.Instance(gce_setup=gce_setup) + request = notebooks_v2.CreateInstanceRequest( + parent=parent, + instance_id=instance_id, + instance=instance, + ) + + try: + operation = client.create_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + # The operation completes once the instance is accepted, before the + # VM finishes provisioning, so the state is usually PROVISIONING here. + response = operation.result(timeout=300) + secure_boot = response.gce_setup.shielded_instance_config.enable_secure_boot + print(f"Resource ID: {response.name}") + print(f"Machine Type: {response.gce_setup.machine_type}") + print(f"Secure Boot Enabled: {secure_boot}") + print(f"State: {response.state.name}") + print( + "The instance is still starting up; call get_instance until the " + "state is ACTIVE before connecting to JupyterLab." + ) + except exceptions.AlreadyExists as err: + print( + f"error: Instance '{instance_id}' already exists in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_create_instance_confidential] diff --git a/notebooks/v2/create_instance_custom_container.py b/notebooks/v2/create_instance_custom_container.py new file mode 100644 index 00000000000..0c6fcf738de --- /dev/null +++ b/notebooks/v2/create_instance_custom_container.py @@ -0,0 +1,75 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_create_instance_custom_container] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def create_instance_custom_container( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-container-instance", + repository: str = "gcr.io/deeplearning-platform-release/base-cpu", +) -> None: + """Creates a Workbench instance from a custom container. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone for the new instance. + instance_id: The ID to assign to the new Workbench instance. + repository: The container image repository URI. + """ + client = notebooks_v2.NotebookServiceClient() + parent = f"projects/{project_id}/locations/{location_id}" + container_image = notebooks_v2.ContainerImage( + repository=repository, + tag="latest", + ) + gce_setup = notebooks_v2.GceSetup( + machine_type="e2-standard-4", + container_image=container_image, + ) + instance = notebooks_v2.Instance(gce_setup=gce_setup) + request = notebooks_v2.CreateInstanceRequest( + parent=parent, + instance_id=instance_id, + instance=instance, + ) + + try: + operation = client.create_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + # The operation completes once the instance is accepted, before the + # VM finishes provisioning, so the state is usually PROVISIONING here. + response = operation.result(timeout=300) + repo = response.gce_setup.container_image.repository + print(f"Resource ID: {response.name}") + print(f"Container Repository: {repo}") + print(f"State: {response.state.name}") + print( + "The instance is still starting up; call get_instance until the " + "state is ACTIVE before connecting to JupyterLab." + ) + except exceptions.AlreadyExists as err: + print( + f"error: Instance '{instance_id}' already exists in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_create_instance_custom_container] diff --git a/notebooks/v2/create_instance_gpu.py b/notebooks/v2/create_instance_gpu.py new file mode 100644 index 00000000000..3b3ebd07950 --- /dev/null +++ b/notebooks/v2/create_instance_gpu.py @@ -0,0 +1,77 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_create_instance_gpu] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def create_instance_gpu( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-gpu-instance", +) -> None: + """Creates a GPU-accelerated Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone for the new GPU instance. + instance_id: The ID to assign to the new Workbench instance. + """ + client = notebooks_v2.NotebookServiceClient() + parent = f"projects/{project_id}/locations/{location_id}" + accelerator_config = notebooks_v2.AcceleratorConfig( + type_=notebooks_v2.AcceleratorConfig.AcceleratorType.NVIDIA_TESLA_T4, + core_count=1, + ) + gpu_driver_config = notebooks_v2.GPUDriverConfig( + enable_gpu_driver=True, + ) + gce_setup = notebooks_v2.GceSetup( + machine_type="n1-standard-4", + accelerator_configs=[accelerator_config], + gpu_driver_config=gpu_driver_config, + ) + instance = notebooks_v2.Instance(gce_setup=gce_setup) + request = notebooks_v2.CreateInstanceRequest( + parent=parent, + instance_id=instance_id, + instance=instance, + ) + + try: + operation = client.create_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + # The operation completes once the instance is accepted, before the + # VM finishes provisioning, so the state is usually PROVISIONING here. + response = operation.result(timeout=300) + print(f"Resource ID: {response.name}") + print(f"Creation Date: {response.create_time}") + print(f"Machine Type: {response.gce_setup.machine_type}") + print(f"State: {response.state.name}") + print( + "The instance is still starting up; call get_instance until the " + "state is ACTIVE before connecting to JupyterLab." + ) + except exceptions.AlreadyExists as err: + print( + f"error: Instance '{instance_id}' already exists in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_create_instance_gpu] diff --git a/notebooks/v2/create_instance_hardened.py b/notebooks/v2/create_instance_hardened.py new file mode 100644 index 00000000000..10e5b4a2c04 --- /dev/null +++ b/notebooks/v2/create_instance_hardened.py @@ -0,0 +1,139 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_create_instance_hardened] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def create_instance_hardened( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-hardened-instance", + network_id: str = "example-network", + subnet_id: str = "example-subnet", + service_account_email: str = "example-workbench-sa@example-project-id.iam.gserviceaccount.com", + instance_owner: str = "example-user@example.com", + key_ring_id: str = "example-key-ring", + crypto_key_id: str = "example-key", +) -> None: + """Creates a security-hardened Workbench instance. + + Combines enterprise hardening controls for regulated environments: + disables public IP, restricts proxy access to a single user with a custom + service account, enables Shielded VM integrity checks, encrypts disks with + CMEK, and disables root access, file downloads, and the JupyterLab terminal + via instance metadata. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone for the new instance. + instance_id: The ID to assign to the new Workbench instance. + network_id: The ID of the VPC network in the project. + subnet_id: The ID of the subnetwork in the instance region. + service_account_email: The email address of the service account to + attach to the instance. + instance_owner: The single user email permitted to access the instance. + key_ring_id: The Cloud KMS key ring ID in the instance region. + crypto_key_id: The Cloud KMS crypto key ID used for disk encryption. + """ + client = notebooks_v2.NotebookServiceClient() + parent = f"projects/{project_id}/locations/{location_id}" + parts = location_id.split("-") + region_id = "-".join(parts[:-1]) if len(parts) > 2 else location_id + network = f"projects/{project_id}/global/networks/{network_id}" + subnet = f"projects/{project_id}/regions/{region_id}/subnetworks/{subnet_id}" + kms_key = ( + f"projects/{project_id}/locations/{region_id}/" + f"keyRings/{key_ring_id}/cryptoKeys/{crypto_key_id}" + ) + network_interface = notebooks_v2.NetworkInterface( + network=network, + subnet=subnet, + nic_type=notebooks_v2.NetworkInterface.NicType.GVNIC, + ) + service_account = notebooks_v2.ServiceAccount( + email=service_account_email, + ) + shielded_config = notebooks_v2.ShieldedInstanceConfig( + enable_secure_boot=True, + enable_vtpm=True, + enable_integrity_monitoring=True, + ) + boot_disk = notebooks_v2.BootDisk( + disk_size_gb=150, + disk_type=notebooks_v2.DiskType.PD_SSD, + disk_encryption=notebooks_v2.DiskEncryption.CMEK, + kms_key=kms_key, + ) + data_disk = notebooks_v2.DataDisk( + disk_size_gb=100, + disk_type=notebooks_v2.DiskType.PD_SSD, + disk_encryption=notebooks_v2.DiskEncryption.CMEK, + kms_key=kms_key, + ) + hardened_metadata = { + "notebook-disable-root": "true", + "notebook-disable-downloads": "true", + "notebook-disable-terminal": "true", + "serial-port-enable": "false", + "block-project-ssh-keys": "true", + } + gce_setup = notebooks_v2.GceSetup( + machine_type="e2-standard-4", + disable_public_ip=True, + network_interfaces=[network_interface], + service_accounts=[service_account], + shielded_instance_config=shielded_config, + boot_disk=boot_disk, + data_disks=[data_disk], + metadata=hardened_metadata, + ) + instance = notebooks_v2.Instance( + gce_setup=gce_setup, + instance_owners=[instance_owner], + ) + request = notebooks_v2.CreateInstanceRequest( + parent=parent, + instance_id=instance_id, + instance=instance, + ) + + try: + operation = client.create_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + # The operation completes once the instance is accepted, before the + # VM finishes provisioning, so the state is usually PROVISIONING here. + response = operation.result(timeout=300) + secure_boot = response.gce_setup.shielded_instance_config.enable_secure_boot + print(f"Resource ID: {response.name}") + print(f"Creation Date: {response.create_time}") + print(f"Disable Public IP: {response.gce_setup.disable_public_ip}") + print(f"Secure Boot Enabled: {secure_boot}") + print(f"State: {response.state.name}") + print( + "The instance is still starting up; call get_instance until the " + "state is ACTIVE before connecting to JupyterLab." + ) + except exceptions.AlreadyExists as err: + print( + f"error: Instance '{instance_id}' already exists in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_create_instance_hardened] diff --git a/notebooks/v2/create_instance_private_vpc.py b/notebooks/v2/create_instance_private_vpc.py new file mode 100644 index 00000000000..56233c60bf0 --- /dev/null +++ b/notebooks/v2/create_instance_private_vpc.py @@ -0,0 +1,88 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_create_instance_private_vpc] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def create_instance_private_vpc( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-private-instance", + network_id: str = "example-network", + subnet_id: str = "example-subnet", +) -> None: + """Creates a Workbench instance with private IP in a custom VPC subnet. + + Disabling the external public IP address isolates the Workbench instance + from the public internet while allowing private connectivity through a + custom VPC network and Cloud NAT or Private Google Access. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone for the new instance. + instance_id: The ID to assign to the new Workbench instance. + network_id: The ID of the VPC network in the project. + subnet_id: The ID of the subnetwork in the instance region. + """ + client = notebooks_v2.NotebookServiceClient() + parent = f"projects/{project_id}/locations/{location_id}" + parts = location_id.split("-") + region_id = "-".join(parts[:-1]) if len(parts) > 2 else location_id + network = f"projects/{project_id}/global/networks/{network_id}" + subnet = f"projects/{project_id}/regions/{region_id}/subnetworks/{subnet_id}" + network_interface = notebooks_v2.NetworkInterface( + network=network, + subnet=subnet, + nic_type=notebooks_v2.NetworkInterface.NicType.GVNIC, + ) + gce_setup = notebooks_v2.GceSetup( + machine_type="e2-standard-4", + disable_public_ip=True, + network_interfaces=[network_interface], + ) + instance = notebooks_v2.Instance(gce_setup=gce_setup) + request = notebooks_v2.CreateInstanceRequest( + parent=parent, + instance_id=instance_id, + instance=instance, + ) + + try: + operation = client.create_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + # The operation completes once the instance is accepted, before the + # VM finishes provisioning, so the state is usually PROVISIONING here. + response = operation.result(timeout=300) + print(f"Resource ID: {response.name}") + print(f"Creation Date: {response.create_time}") + print(f"Disable Public IP: {response.gce_setup.disable_public_ip}") + print(f"Machine Type: {response.gce_setup.machine_type}") + print(f"State: {response.state.name}") + print( + "The instance is still starting up; call get_instance until the " + "state is ACTIVE before connecting to JupyterLab." + ) + except exceptions.AlreadyExists as err: + print( + f"error: Instance '{instance_id}' already exists in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_create_instance_private_vpc] diff --git a/notebooks/v2/create_instance_single_user.py b/notebooks/v2/create_instance_single_user.py new file mode 100644 index 00000000000..3b949dc5952 --- /dev/null +++ b/notebooks/v2/create_instance_single_user.py @@ -0,0 +1,77 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_create_instance_single_user] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def create_instance_single_user( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-single-user-instance", + instance_owner: str = "example-user@example.com", + service_account_email: str = "example-sa@example-project-id.iam.gserviceaccount.com", +) -> None: + """Creates a single-user Workbench instance with a custom service account. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone for the new instance. + instance_id: The ID to assign to the new Workbench instance. + instance_owner: The email address of the single permitted user. + service_account_email: The email address of the service account to + attach to the instance. + """ + client = notebooks_v2.NotebookServiceClient() + parent = f"projects/{project_id}/locations/{location_id}" + service_account = notebooks_v2.ServiceAccount(email=service_account_email) + gce_setup = notebooks_v2.GceSetup( + machine_type="e2-standard-4", + service_accounts=[service_account], + ) + instance = notebooks_v2.Instance( + gce_setup=gce_setup, + instance_owners=[instance_owner], + ) + request = notebooks_v2.CreateInstanceRequest( + parent=parent, + instance_id=instance_id, + instance=instance, + ) + + try: + operation = client.create_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + # The operation completes once the instance is accepted, before the + # VM finishes provisioning, so the state is usually PROVISIONING here. + response = operation.result(timeout=300) + print(f"Resource ID: {response.name}") + print(f"Creator: {response.creator}") + print(f"State: {response.state.name}") + print( + "The instance is still starting up; call get_instance until the " + "state is ACTIVE before connecting to JupyterLab." + ) + except exceptions.AlreadyExists as err: + print( + f"error: Instance '{instance_id}' already exists in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_create_instance_single_user] diff --git a/notebooks/v2/delete_instance.py b/notebooks/v2/delete_instance.py new file mode 100644 index 00000000000..82988dab1a7 --- /dev/null +++ b/notebooks/v2/delete_instance.py @@ -0,0 +1,53 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_delete_instance] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def delete_instance( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Deletes a Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to delete. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + request = notebooks_v2.DeleteInstanceRequest(name=name) + + try: + operation = client.delete_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + operation.result(timeout=300) + print(f"Resource ID: {name}") + print(f"Location: {location_id}") + print(f"Operation Done: {operation.done()}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_delete_instance] diff --git a/notebooks/v2/diagnose_instance.py b/notebooks/v2/diagnose_instance.py new file mode 100644 index 00000000000..e1720305d44 --- /dev/null +++ b/notebooks/v2/diagnose_instance.py @@ -0,0 +1,63 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_diagnose_instance] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def diagnose_instance( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", + gcs_bucket: str = "gs://example-bucket", +) -> None: + """Runs diagnostics on a Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to diagnose. + gcs_bucket: The Cloud Storage bucket URI to store diagnostic logs. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + diagnostic_config = notebooks_v2.DiagnosticConfig( + gcs_bucket=gcs_bucket, + relative_path="example-diagnostics", + enable_repair_flag=False, + ) + request = notebooks_v2.DiagnoseInstanceRequest( + name=name, + diagnostic_config=diagnostic_config, + ) + + try: + operation = client.diagnose_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result(timeout=600) + print(f"Resource ID: {response.name}") + print(f"Health State: {response.health_state.name}") + print(f"State: {response.state.name}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_diagnose_instance] diff --git a/notebooks/v2/get_instance.py b/notebooks/v2/get_instance.py new file mode 100644 index 00000000000..f09c36d1b3d --- /dev/null +++ b/notebooks/v2/get_instance.py @@ -0,0 +1,53 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_get_instance] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def get_instance( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Gets details for a Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to retrieve. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + request = notebooks_v2.GetInstanceRequest(name=name) + + try: + response = client.get_instance(request=request) + print(f"Resource ID: {response.name}") + print(f"Creation Date: {response.create_time}") + print(f"Health State: {response.health_state.name}") + print(f"Proxy URI: {response.proxy_uri}") + print(f"State: {response.state.name}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_get_instance] diff --git a/notebooks/v2/get_instance_iam_policy.py b/notebooks/v2/get_instance_iam_policy.py new file mode 100644 index 00000000000..b42559951e8 --- /dev/null +++ b/notebooks/v2/get_instance_iam_policy.py @@ -0,0 +1,57 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_get_instance_iam_policy] +import base64 +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +from google.iam.v1 import iam_policy_pb2 + + +def get_instance_iam_policy( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Gets the IAM policy for a Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + request = iam_policy_pb2.GetIamPolicyRequest(resource=name) + + try: + response = client.get_iam_policy(request=request) + print(f"Resource ID: {name}") + # The etag is an opaque byte string; base64 matches how gcloud and + # the REST API display it. + print(f"Policy Etag: {base64.b64encode(response.etag).decode()}") + print(f"Policy Version: {response.version}") + for binding in response.bindings: + print(f"Role: {binding.role}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_get_instance_iam_policy] diff --git a/notebooks/v2/list_instances.py b/notebooks/v2/list_instances.py new file mode 100644 index 00000000000..d00ff9ba1ec --- /dev/null +++ b/notebooks/v2/list_instances.py @@ -0,0 +1,52 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_list_instances] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def list_instances( + project_id: str, + location_id: str = "us-central1-a", +) -> None: + """Lists Workbench instances in a zone. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone to list instances from. + """ + client = notebooks_v2.NotebookServiceClient() + parent = f"projects/{project_id}/locations/{location_id}" + request = notebooks_v2.ListInstancesRequest(parent=parent) + + try: + response = client.list_instances(request=request) + for instance in response: + print(f"Resource ID: {instance.name}") + print(f"Creation Date: {instance.create_time}") + print(f"Instance ID: {instance.id}") + print(f"Proxy URI: {instance.proxy_uri}") + print(f"State: {instance.state.name}") + except exceptions.NotFound as err: + print( + f"error: Could not list instances in location '{location_id}' " + f"for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_list_instances] diff --git a/notebooks/v2/noxfile_config.py b/notebooks/v2/noxfile_config.py new file mode 100644 index 00000000000..cb3ff7a91a6 --- /dev/null +++ b/notebooks/v2/noxfile_config.py @@ -0,0 +1,31 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Default TEST_CONFIG_OVERRIDE for python repos. + +TEST_CONFIG_OVERRIDE = { + # You can opt out from the test for specific Python versions. + "ignored_versions": ["3.8", "3.9", "3.11", "3.12", "3.13"], + # Old samples are opted out of enforcing Python type hints + # All new samples should feature them + "enforce_type_hints": True, + # An envvar key for determining the project id to use. + "gcloud_project_env": "GOOGLE_CLOUD_PROJECT", + # If you need to use a specific version of pip, + # change pip_version_override to the string representation + # of the version number, for example, "20.2.4" + "pip_version_override": None, + # A dictionary you want to inject into your test. + "envs": {}, +} diff --git a/notebooks/v2/requirements-test.txt b/notebooks/v2/requirements-test.txt new file mode 100644 index 00000000000..b3b2b3ca4cc --- /dev/null +++ b/notebooks/v2/requirements-test.txt @@ -0,0 +1 @@ +pytest==9.0.3 diff --git a/notebooks/v2/requirements.txt b/notebooks/v2/requirements.txt new file mode 100644 index 00000000000..877976ff7b1 --- /dev/null +++ b/notebooks/v2/requirements.txt @@ -0,0 +1 @@ +google-cloud-notebooks==1.18.0 diff --git a/notebooks/v2/reset_instance.py b/notebooks/v2/reset_instance.py new file mode 100644 index 00000000000..056c54827c2 --- /dev/null +++ b/notebooks/v2/reset_instance.py @@ -0,0 +1,52 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_reset_instance] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def reset_instance( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Resets a Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to reset. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + request = notebooks_v2.ResetInstanceRequest(name=name) + + try: + operation = client.reset_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result(timeout=300) + print(f"Resource ID: {response.name}") + print(f"State: {response.state.name}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_reset_instance] diff --git a/notebooks/v2/rollback_instance.py b/notebooks/v2/rollback_instance.py new file mode 100644 index 00000000000..c73a61d1ffa --- /dev/null +++ b/notebooks/v2/rollback_instance.py @@ -0,0 +1,61 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_rollback_instance] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def rollback_instance( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", + snapshot_id: str = "example-snapshot", + revision_id: str = "1", +) -> None: + """Rolls back a Workbench instance to a snapshot. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to roll back. + snapshot_id: The ID of the Compute Engine snapshot in the project. + revision_id: The revision ID to roll back to. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + target_snapshot = f"projects/{project_id}/global/snapshots/{snapshot_id}" + request = notebooks_v2.RollbackInstanceRequest( + name=name, + target_snapshot=target_snapshot, + revision_id=revision_id, + ) + + try: + operation = client.rollback_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result(timeout=600) + print(f"Resource ID: {response.name}") + print(f"State: {response.state.name}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_rollback_instance] diff --git a/notebooks/v2/set_instance_iam_policy.py b/notebooks/v2/set_instance_iam_policy.py new file mode 100644 index 00000000000..3ad73155b2f --- /dev/null +++ b/notebooks/v2/set_instance_iam_policy.py @@ -0,0 +1,75 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_set_instance_iam_policy] +import base64 +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +from google.iam.v1 import iam_policy_pb2 +from google.iam.v1 import policy_pb2 + + +def set_instance_iam_policy( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", + member: str = "user:example-user@example.com", + role: str = "roles/notebooks.runner", +) -> None: + """Sets the IAM policy for a Workbench instance using read-modify-write. + + Retrieves the existing IAM policy first so that appending the new role + binding preserves existing access controls and optimistic concurrency etags. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance. + member: The principal to grant the role to. + role: The IAM role to bind on the instance. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + get_request = iam_policy_pb2.GetIamPolicyRequest(resource=name) + binding = policy_pb2.Binding( + role=role, + members=[member], + ) + + try: + # Read the current IAM policy first to preserve existing bindings and etag. + policy = client.get_iam_policy(request=get_request) + policy.bindings.append(binding) + set_request = iam_policy_pb2.SetIamPolicyRequest( + resource=name, + policy=policy, + ) + response = client.set_iam_policy(request=set_request) + print(f"Resource ID: {name}") + print(f"Bindings Count: {len(response.bindings)}") + # The etag is an opaque byte string; base64 matches how gcloud and + # the REST API display it. + print(f"Policy Etag: {base64.b64encode(response.etag).decode()}") + print(f"Policy Version: {response.version}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_set_instance_iam_policy] diff --git a/notebooks/v2/start_instance.py b/notebooks/v2/start_instance.py new file mode 100644 index 00000000000..177bb545ded --- /dev/null +++ b/notebooks/v2/start_instance.py @@ -0,0 +1,59 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_start_instance] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def start_instance( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Starts a stopped Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to start. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + request = notebooks_v2.StartInstanceRequest(name=name) + + try: + operation = client.start_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + # The operation completes before the VM finishes booting, so the + # state is usually PROVISIONING here. + response = operation.result(timeout=300) + print(f"Resource ID: {response.name}") + print(f"Proxy URI: {response.proxy_uri}") + print(f"State: {response.state.name}") + print( + "The instance is still starting up; call get_instance until the " + "state is ACTIVE before connecting to JupyterLab." + ) + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_start_instance] diff --git a/notebooks/v2/stop_instance.py b/notebooks/v2/stop_instance.py new file mode 100644 index 00000000000..e05415e3421 --- /dev/null +++ b/notebooks/v2/stop_instance.py @@ -0,0 +1,52 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_stop_instance] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def stop_instance( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Stops a running Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to stop. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + request = notebooks_v2.StopInstanceRequest(name=name) + + try: + operation = client.stop_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result(timeout=300) + print(f"Resource ID: {response.name}") + print(f"State: {response.state.name}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_stop_instance] diff --git a/notebooks/v2/tests/__init__.py b/notebooks/v2/tests/__init__.py new file mode 100644 index 00000000000..58d482ea386 --- /dev/null +++ b/notebooks/v2/tests/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/notebooks/v2/tests/check_instance_iam_permissions_test.py b/notebooks/v2/tests/check_instance_iam_permissions_test.py new file mode 100644 index 00000000000..797128f409e --- /dev/null +++ b/notebooks/v2/tests/check_instance_iam_permissions_test.py @@ -0,0 +1,72 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +from google.iam.v1 import iam_policy_pb2 +import pytest + +from .. import check_instance_iam_permissions + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_check_instance_iam_permissions_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.test_iam_permissions.return_value = ( + iam_policy_pb2.TestIamPermissionsResponse( + permissions=["notebooks.instances.get"] + ) + ) + + check_instance_iam_permissions.check_instance_iam_permissions( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + request = mock_client.test_iam_permissions.call_args.kwargs["request"] + assert request.resource == INSTANCE_NAME + assert list(request.permissions) == ["notebooks.instances.get"] + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Allowed Permissions: ['notebooks.instances.get']" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_check_instance_iam_permissions_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.test_iam_permissions.side_effect = exceptions.NotFound( + "Instance not found" + ) + + check_instance_iam_permissions.check_instance_iam_permissions( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/check_instance_upgradability_test.py b/notebooks/v2/tests/check_instance_upgradability_test.py new file mode 100644 index 00000000000..1eb4919c3dd --- /dev/null +++ b/notebooks/v2/tests/check_instance_upgradability_test.py @@ -0,0 +1,73 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import check_instance_upgradability + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_check_instance_upgradability_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.check_instance_upgradability.return_value = ( + notebooks_v2.CheckInstanceUpgradabilityResponse( + upgradeable=True, + upgrade_version="m120", + upgrade_info="Upgrade available", + upgrade_image="gcr.io/deeplearning-platform-release/base-cpu:m120", + ) + ) + + check_instance_upgradability.check_instance_upgradability( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + request = mock_client.check_instance_upgradability.call_args.kwargs["request"] + assert request.notebook_instance == INSTANCE_NAME + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Upgradeable: True" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_check_instance_upgradability_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.check_instance_upgradability.side_effect = exceptions.NotFound( + "Instance not found" + ) + + check_instance_upgradability.check_instance_upgradability( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/create_instance_cmek_test.py b/notebooks/v2/tests/create_instance_cmek_test.py new file mode 100644 index 00000000000..ce28710498d --- /dev/null +++ b/notebooks/v2/tests/create_instance_cmek_test.py @@ -0,0 +1,112 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import create_instance_cmek + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-cmek-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_create_instance_cmek_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + gce_setup=notebooks_v2.GceSetup( + boot_disk=notebooks_v2.BootDisk( + disk_encryption=notebooks_v2.DiskEncryption.CMEK, + kms_key=( + f"projects/{PROJECT_ID}/locations/us-central1/" + "keyRings/example-key-ring/cryptoKeys/example-key" + ), + ) + ), + ) + mock_client.create_instance.return_value = operation + + create_instance_cmek.create_instance_cmek(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + request = mock_client.create_instance.call_args.kwargs["request"] + region_id = "-".join(LOCATION_ID.split("-")[:-1]) + expected_kms_key = ( + f"projects/{PROJECT_ID}/locations/{region_id}/" + "keyRings/example-key-ring/cryptoKeys/example-key" + ) + actual_request = notebooks_v2.CreateInstanceRequest.to_dict( + request, + always_print_fields_with_no_presence=False, + use_integers_for_enums=False, + ) + assert actual_request == { + "parent": f"projects/{PROJECT_ID}/locations/{LOCATION_ID}", + "instance_id": INSTANCE_ID, + "instance": { + "gce_setup": { + "machine_type": "e2-standard-4", + "boot_disk": { + "disk_size_gb": "150", + "disk_type": "PD_SSD", + "disk_encryption": "CMEK", + "kms_key": expected_kms_key, + }, + "data_disks": [ + { + "disk_size_gb": "100", + "disk_type": "PD_SSD", + "disk_encryption": "CMEK", + "kms_key": expected_kms_key, + } + ], + }, + }, + } + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Boot Disk Encryption: CMEK" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_create_instance_cmek_already_exists( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.create_instance.side_effect = exceptions.AlreadyExists( + "Instance already exists" + ) + + create_instance_cmek.create_instance_cmek(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/create_instance_confidential_test.py b/notebooks/v2/tests/create_instance_confidential_test.py new file mode 100644 index 00000000000..a3b7c8a32b6 --- /dev/null +++ b/notebooks/v2/tests/create_instance_confidential_test.py @@ -0,0 +1,89 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import create_instance_confidential + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-confidential-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_create_instance_confidential_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + gce_setup=notebooks_v2.GceSetup( + machine_type="n2d-standard-4", + shielded_instance_config=notebooks_v2.ShieldedInstanceConfig( + enable_secure_boot=True, + enable_vtpm=True, + enable_integrity_monitoring=True, + ), + ), + ) + mock_client.create_instance.return_value = operation + + create_instance_confidential.create_instance_confidential( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + request = mock_client.create_instance.call_args.kwargs["request"] + assert request.parent == f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" + assert request.instance_id == INSTANCE_ID + gce_setup = request.instance.gce_setup + shielded = gce_setup.shielded_instance_config + assert shielded.enable_secure_boot + assert shielded.enable_vtpm + assert shielded.enable_integrity_monitoring + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Secure Boot Enabled: True" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_create_instance_confidential_already_exists( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.create_instance.side_effect = exceptions.AlreadyExists( + "Instance already exists" + ) + + create_instance_confidential.create_instance_confidential( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/create_instance_custom_container_test.py b/notebooks/v2/tests/create_instance_custom_container_test.py new file mode 100644 index 00000000000..9a172d133c0 --- /dev/null +++ b/notebooks/v2/tests/create_instance_custom_container_test.py @@ -0,0 +1,90 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import create_instance_custom_container + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-container-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_create_instance_custom_container_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + gce_setup=notebooks_v2.GceSetup( + container_image=notebooks_v2.ContainerImage( + repository="gcr.io/deeplearning-platform-release/base-cpu", + tag="latest", + ) + ), + ) + mock_client.create_instance.return_value = operation + + create_instance_custom_container.create_instance_custom_container( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + request = mock_client.create_instance.call_args.kwargs["request"] + assert request.parent == f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" + assert request.instance_id == INSTANCE_ID + gce_setup = request.instance.gce_setup + assert gce_setup.container_image.repository == ( + "gcr.io/deeplearning-platform-release/base-cpu" + ) + assert gce_setup.container_image.tag == "latest" + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert ( + "Container Repository: gcr.io/deeplearning-platform-release/base-cpu" + in captured.out + ) + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_create_instance_custom_container_already_exists( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.create_instance.side_effect = exceptions.AlreadyExists( + "Instance already exists" + ) + + create_instance_custom_container.create_instance_custom_container( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/create_instance_gpu_test.py b/notebooks/v2/tests/create_instance_gpu_test.py new file mode 100644 index 00000000000..3da54cb01df --- /dev/null +++ b/notebooks/v2/tests/create_instance_gpu_test.py @@ -0,0 +1,80 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import create_instance_gpu + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-gpu-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_create_instance_gpu_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + gce_setup=notebooks_v2.GceSetup(machine_type="n1-standard-4"), + ) + mock_client.create_instance.return_value = operation + + create_instance_gpu.create_instance_gpu(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + request = mock_client.create_instance.call_args.kwargs["request"] + assert request.parent == f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" + assert request.instance_id == INSTANCE_ID + gce_setup = request.instance.gce_setup + accelerator = gce_setup.accelerator_configs[0] + assert accelerator.type_ == ( + notebooks_v2.AcceleratorConfig.AcceleratorType.NVIDIA_TESLA_T4 + ) + assert accelerator.core_count == 1 + assert gce_setup.gpu_driver_config.enable_gpu_driver + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Machine Type: n1-standard-4" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_create_instance_gpu_already_exists( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.create_instance.side_effect = exceptions.AlreadyExists( + "Instance already exists" + ) + + create_instance_gpu.create_instance_gpu(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/create_instance_hardened_test.py b/notebooks/v2/tests/create_instance_hardened_test.py new file mode 100644 index 00000000000..61bcef0ae4d --- /dev/null +++ b/notebooks/v2/tests/create_instance_hardened_test.py @@ -0,0 +1,150 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import create_instance_hardened + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-hardened-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_create_instance_hardened_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + gce_setup=notebooks_v2.GceSetup( + disable_public_ip=True, + boot_disk=notebooks_v2.BootDisk( + disk_encryption=notebooks_v2.DiskEncryption.CMEK + ), + shielded_instance_config=notebooks_v2.ShieldedInstanceConfig( + enable_secure_boot=True, + enable_vtpm=True, + enable_integrity_monitoring=True, + ), + ), + ) + mock_client.create_instance.return_value = operation + + custom_sa_email = "custom-workbench-sa@shared-project.iam.gserviceaccount.com" + create_instance_hardened.create_instance_hardened( + PROJECT_ID, + LOCATION_ID, + INSTANCE_ID, + service_account_email=custom_sa_email, + ) + + request = mock_client.create_instance.call_args.kwargs["request"] + region_id = "-".join(LOCATION_ID.split("-")[:-1]) + expected_kms_key = ( + f"projects/{PROJECT_ID}/locations/{region_id}/" + "keyRings/example-key-ring/cryptoKeys/example-key" + ) + actual_request = notebooks_v2.CreateInstanceRequest.to_dict( + request, + always_print_fields_with_no_presence=False, + use_integers_for_enums=False, + ) + assert actual_request == { + "parent": f"projects/{PROJECT_ID}/locations/{LOCATION_ID}", + "instance_id": INSTANCE_ID, + "instance": { + "instance_owners": ["example-user@example.com"], + "gce_setup": { + "machine_type": "e2-standard-4", + "disable_public_ip": True, + "service_accounts": [{"email": custom_sa_email}], + "boot_disk": { + "disk_size_gb": "150", + "disk_type": "PD_SSD", + "disk_encryption": "CMEK", + "kms_key": expected_kms_key, + }, + "data_disks": [ + { + "disk_size_gb": "100", + "disk_type": "PD_SSD", + "disk_encryption": "CMEK", + "kms_key": expected_kms_key, + } + ], + "network_interfaces": [ + { + "network": ( + f"projects/{PROJECT_ID}/global/networks/example-network" + ), + "subnet": ( + f"projects/{PROJECT_ID}/regions/{region_id}/" + "subnetworks/example-subnet" + ), + "nic_type": "GVNIC", + } + ], + "shielded_instance_config": { + "enable_secure_boot": True, + "enable_vtpm": True, + "enable_integrity_monitoring": True, + }, + "metadata": { + "notebook-disable-root": "true", + "notebook-disable-downloads": "true", + "notebook-disable-terminal": "true", + "serial-port-enable": "false", + "block-project-ssh-keys": "true", + }, + }, + }, + } + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Disable Public IP: True" in captured.out + assert "Secure Boot Enabled: True" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_create_instance_hardened_already_exists( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.create_instance.side_effect = exceptions.AlreadyExists( + "Instance already exists" + ) + + create_instance_hardened.create_instance_hardened( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/create_instance_private_vpc_test.py b/notebooks/v2/tests/create_instance_private_vpc_test.py new file mode 100644 index 00000000000..5de23bd2277 --- /dev/null +++ b/notebooks/v2/tests/create_instance_private_vpc_test.py @@ -0,0 +1,106 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import create_instance_private_vpc + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-private-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_create_instance_private_vpc_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + gce_setup=notebooks_v2.GceSetup( + machine_type="e2-standard-4", + disable_public_ip=True, + ), + ) + mock_client.create_instance.return_value = operation + + create_instance_private_vpc.create_instance_private_vpc( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + request = mock_client.create_instance.call_args.kwargs["request"] + region_id = "-".join(LOCATION_ID.split("-")[:-1]) + actual_request = notebooks_v2.CreateInstanceRequest.to_dict( + request, + always_print_fields_with_no_presence=False, + use_integers_for_enums=False, + ) + assert actual_request == { + "parent": f"projects/{PROJECT_ID}/locations/{LOCATION_ID}", + "instance_id": INSTANCE_ID, + "instance": { + "gce_setup": { + "machine_type": "e2-standard-4", + "disable_public_ip": True, + "network_interfaces": [ + { + "network": ( + f"projects/{PROJECT_ID}/global/networks/example-network" + ), + "subnet": ( + f"projects/{PROJECT_ID}/regions/{region_id}/" + "subnetworks/example-subnet" + ), + "nic_type": "GVNIC", + } + ], + }, + }, + } + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Disable Public IP: True" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_create_instance_private_vpc_already_exists( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.create_instance.side_effect = exceptions.AlreadyExists( + "Instance already exists" + ) + + create_instance_private_vpc.create_instance_private_vpc( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/create_instance_single_user_test.py b/notebooks/v2/tests/create_instance_single_user_test.py new file mode 100644 index 00000000000..521c503691f --- /dev/null +++ b/notebooks/v2/tests/create_instance_single_user_test.py @@ -0,0 +1,85 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import create_instance_single_user + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-single-user-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_create_instance_single_user_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + creator="example-user@example.com", + state=notebooks_v2.State.ACTIVE, + gce_setup=notebooks_v2.GceSetup(machine_type="e2-standard-4"), + ) + mock_client.create_instance.return_value = operation + + custom_sa_email = "123456789012-compute@developer.gserviceaccount.com" + create_instance_single_user.create_instance_single_user( + PROJECT_ID, + LOCATION_ID, + INSTANCE_ID, + service_account_email=custom_sa_email, + ) + + request = mock_client.create_instance.call_args.kwargs["request"] + assert request.parent == f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" + assert request.instance_id == INSTANCE_ID + gce_setup = request.instance.gce_setup + assert list(request.instance.instance_owners) == ["example-user@example.com"] + assert gce_setup.service_accounts[0].email == custom_sa_email + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Creator: example-user@example.com" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_create_instance_single_user_already_exists( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.create_instance.side_effect = exceptions.AlreadyExists( + "Instance already exists" + ) + + create_instance_single_user.create_instance_single_user( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/create_instance_test.py b/notebooks/v2/tests/create_instance_test.py new file mode 100644 index 00000000000..e9d77c2f7ee --- /dev/null +++ b/notebooks/v2/tests/create_instance_test.py @@ -0,0 +1,77 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import create_instance + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_create_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + proxy_uri="example-proxy.notebooks.googleusercontent.com", + gce_setup=notebooks_v2.GceSetup(machine_type="e2-standard-4"), + ) + mock_client.create_instance.return_value = operation + + create_instance.create_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + request = mock_client.create_instance.call_args.kwargs["request"] + assert request.parent == f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" + assert request.instance_id == INSTANCE_ID + gce_setup = request.instance.gce_setup + assert gce_setup.machine_type == "e2-standard-4" + + captured = capsys.readouterr() + assert "Waiting for operation to complete:" in captured.out + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Machine Type: e2-standard-4" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_create_instance_already_exists( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.create_instance.side_effect = exceptions.AlreadyExists( + "Instance already exists" + ) + + create_instance.create_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/delete_instance_test.py b/notebooks/v2/tests/delete_instance_test.py new file mode 100644 index 00000000000..66192e9e9d8 --- /dev/null +++ b/notebooks/v2/tests/delete_instance_test.py @@ -0,0 +1,67 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import delete_instance + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_delete_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = None + operation.done.return_value = True + mock_client.delete_instance.return_value = operation + + delete_instance.delete_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + request = mock_client.delete_instance.call_args.kwargs["request"] + assert request.name == INSTANCE_NAME + + captured = capsys.readouterr() + assert "Waiting for operation to complete:" in captured.out + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Operation Done: True" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_delete_instance_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.delete_instance.side_effect = exceptions.NotFound("Instance not found") + + delete_instance.delete_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/diagnose_instance_test.py b/notebooks/v2/tests/diagnose_instance_test.py new file mode 100644 index 00000000000..6388cc1cf95 --- /dev/null +++ b/notebooks/v2/tests/diagnose_instance_test.py @@ -0,0 +1,73 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import diagnose_instance + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_diagnose_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + health_state=notebooks_v2.HealthState.HEALTHY, + state=notebooks_v2.State.ACTIVE, + ) + mock_client.diagnose_instance.return_value = operation + + diagnose_instance.diagnose_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + request = mock_client.diagnose_instance.call_args.kwargs["request"] + assert request.name == INSTANCE_NAME + assert request.diagnostic_config.gcs_bucket == "gs://example-bucket" + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Health State: HEALTHY" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_diagnose_instance_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.diagnose_instance.side_effect = exceptions.NotFound( + "Instance not found" + ) + + diagnose_instance.diagnose_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/get_instance_iam_policy_test.py b/notebooks/v2/tests/get_instance_iam_policy_test.py new file mode 100644 index 00000000000..edbc0e5ca69 --- /dev/null +++ b/notebooks/v2/tests/get_instance_iam_policy_test.py @@ -0,0 +1,75 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +from google.iam.v1 import policy_pb2 +import pytest + +from .. import get_instance_iam_policy + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_get_instance_iam_policy_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.get_iam_policy.return_value = policy_pb2.Policy( + version=1, + etag=b"etag-1", + bindings=[ + policy_pb2.Binding( + role="roles/notebooks.runner", + members=["user:example-user@example.com"], + ) + ], + ) + + get_instance_iam_policy.get_instance_iam_policy( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + request = mock_client.get_iam_policy.call_args.kwargs["request"] + assert request.resource == INSTANCE_NAME + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Policy Etag: ZXRhZy0x" in captured.out + assert "Role: roles/notebooks.runner" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_get_instance_iam_policy_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.get_iam_policy.side_effect = exceptions.NotFound("Instance not found") + + get_instance_iam_policy.get_instance_iam_policy( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/get_instance_test.py b/notebooks/v2/tests/get_instance_test.py new file mode 100644 index 00000000000..74f86ea9b29 --- /dev/null +++ b/notebooks/v2/tests/get_instance_test.py @@ -0,0 +1,69 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import get_instance + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_get_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.get_instance.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + id="1234567890", + creator="example-user@example.com", + proxy_uri="example-proxy.notebooks.googleusercontent.com", + state=notebooks_v2.State.ACTIVE, + health_state=notebooks_v2.HealthState.HEALTHY, + gce_setup=notebooks_v2.GceSetup(machine_type="e2-standard-4"), + ) + + get_instance.get_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + request = mock_client.get_instance.call_args.kwargs["request"] + assert request.name == INSTANCE_NAME + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Health State: HEALTHY" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_get_instance_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.get_instance.side_effect = exceptions.NotFound("Instance not found") + + get_instance.get_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/list_instances_test.py b/notebooks/v2/tests/list_instances_test.py new file mode 100644 index 00000000000..d651ca58dc0 --- /dev/null +++ b/notebooks/v2/tests/list_instances_test.py @@ -0,0 +1,68 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import list_instances + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_list_instances_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.list_instances.return_value = [ + notebooks_v2.Instance( + name=INSTANCE_NAME, + id="1234567890", + creator="example-user@example.com", + proxy_uri="example-proxy.notebooks.googleusercontent.com", + state=notebooks_v2.State.ACTIVE, + ) + ] + + list_instances.list_instances(PROJECT_ID, LOCATION_ID) + + request = mock_client.list_instances.call_args.kwargs["request"] + assert request.parent == f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_list_instances_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.list_instances.side_effect = exceptions.NotFound("Zone not found") + + list_instances.list_instances(PROJECT_ID, LOCATION_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert LOCATION_ID in captured.err diff --git a/notebooks/v2/tests/reset_instance_test.py b/notebooks/v2/tests/reset_instance_test.py new file mode 100644 index 00000000000..3b99a450325 --- /dev/null +++ b/notebooks/v2/tests/reset_instance_test.py @@ -0,0 +1,69 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import reset_instance + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_reset_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + ) + mock_client.reset_instance.return_value = operation + + reset_instance.reset_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + request = mock_client.reset_instance.call_args.kwargs["request"] + assert request.name == INSTANCE_NAME + + captured = capsys.readouterr() + assert "Waiting for operation to complete:" in captured.out + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_reset_instance_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.reset_instance.side_effect = exceptions.NotFound("Instance not found") + + reset_instance.reset_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/rollback_instance_test.py b/notebooks/v2/tests/rollback_instance_test.py new file mode 100644 index 00000000000..ded7cd16a40 --- /dev/null +++ b/notebooks/v2/tests/rollback_instance_test.py @@ -0,0 +1,74 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import rollback_instance + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_rollback_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + ) + mock_client.rollback_instance.return_value = operation + + rollback_instance.rollback_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + request = mock_client.rollback_instance.call_args.kwargs["request"] + assert request.name == INSTANCE_NAME + assert request.target_snapshot == ( + f"projects/{PROJECT_ID}/global/snapshots/example-snapshot" + ) + assert request.revision_id == "1" + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_rollback_instance_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.rollback_instance.side_effect = exceptions.NotFound( + "Instance not found" + ) + + rollback_instance.rollback_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/set_instance_iam_policy_test.py b/notebooks/v2/tests/set_instance_iam_policy_test.py new file mode 100644 index 00000000000..bf1a2e5ebbc --- /dev/null +++ b/notebooks/v2/tests/set_instance_iam_policy_test.py @@ -0,0 +1,85 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +from google.iam.v1 import policy_pb2 +import pytest + +from .. import set_instance_iam_policy + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_set_instance_iam_policy_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + existing_binding = policy_pb2.Binding( + role="roles/notebooks.admin", + members=["user:existing-admin@example.com"], + ) + mock_client.get_iam_policy.return_value = policy_pb2.Policy( + version=1, + etag=b"etag-1", + bindings=[existing_binding], + ) + mock_client.set_iam_policy.side_effect = lambda request: request.policy + + set_instance_iam_policy.set_instance_iam_policy( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + get_request = mock_client.get_iam_policy.call_args.kwargs["request"] + assert get_request.resource == INSTANCE_NAME + set_request = mock_client.set_iam_policy.call_args.kwargs["request"] + assert set_request.resource == INSTANCE_NAME + # The etag from the read must be sent back so concurrent edits are rejected. + assert set_request.policy.etag == b"etag-1" + bindings = {(b.role, tuple(b.members)) for b in set_request.policy.bindings} + assert bindings == { + ("roles/notebooks.admin", ("user:existing-admin@example.com",)), + ("roles/notebooks.runner", ("user:example-user@example.com",)), + } + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Policy Etag: ZXRhZy0x" in captured.out + assert "Bindings Count: 2" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_set_instance_iam_policy_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.get_iam_policy.side_effect = exceptions.NotFound("Instance not found") + + set_instance_iam_policy.set_instance_iam_policy( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + mock_client.set_iam_policy.assert_not_called() + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/start_instance_test.py b/notebooks/v2/tests/start_instance_test.py new file mode 100644 index 00000000000..0f8e0cb10a2 --- /dev/null +++ b/notebooks/v2/tests/start_instance_test.py @@ -0,0 +1,70 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import start_instance + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_start_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + proxy_uri="example-proxy.notebooks.googleusercontent.com", + state=notebooks_v2.State.ACTIVE, + ) + mock_client.start_instance.return_value = operation + + start_instance.start_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + request = mock_client.start_instance.call_args.kwargs["request"] + assert request.name == INSTANCE_NAME + + captured = capsys.readouterr() + assert "Waiting for operation to complete:" in captured.out + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_start_instance_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.start_instance.side_effect = exceptions.NotFound("Instance not found") + + start_instance.start_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/stop_instance_test.py b/notebooks/v2/tests/stop_instance_test.py new file mode 100644 index 00000000000..8b66c0dd0b8 --- /dev/null +++ b/notebooks/v2/tests/stop_instance_test.py @@ -0,0 +1,69 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import stop_instance + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_stop_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.STOPPED, + ) + mock_client.stop_instance.return_value = operation + + stop_instance.stop_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + request = mock_client.stop_instance.call_args.kwargs["request"] + assert request.name == INSTANCE_NAME + + captured = capsys.readouterr() + assert "Waiting for operation to complete:" in captured.out + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "State: STOPPED" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_stop_instance_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.stop_instance.side_effect = exceptions.NotFound("Instance not found") + + stop_instance.stop_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/update_instance_machine_type_test.py b/notebooks/v2/tests/update_instance_machine_type_test.py new file mode 100644 index 00000000000..a2bdb98b1cb --- /dev/null +++ b/notebooks/v2/tests/update_instance_machine_type_test.py @@ -0,0 +1,76 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import update_instance_machine_type + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_update_instance_machine_type_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.STOPPED, + gce_setup=notebooks_v2.GceSetup(machine_type="n1-standard-8"), + ) + mock_client.update_instance.return_value = operation + + update_instance_machine_type.update_instance_machine_type( + PROJECT_ID, LOCATION_ID, INSTANCE_ID, "n1-standard-8" + ) + + request = mock_client.update_instance.call_args.kwargs["request"] + assert request.instance.name == INSTANCE_NAME + assert request.instance.gce_setup.machine_type == "n1-standard-8" + assert list(request.update_mask.paths) == ["gce_setup.machine_type"] + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Machine Type: n1-standard-8" in captured.out + assert "State: STOPPED" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_update_instance_machine_type_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.update_instance.side_effect = exceptions.NotFound("Instance not found") + + update_instance_machine_type.update_instance_machine_type( + PROJECT_ID, LOCATION_ID, INSTANCE_ID, "n1-standard-8" + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/update_instance_metadata_test.py b/notebooks/v2/tests/update_instance_metadata_test.py new file mode 100644 index 00000000000..5ca1d578ead --- /dev/null +++ b/notebooks/v2/tests/update_instance_metadata_test.py @@ -0,0 +1,84 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import update_instance_metadata + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_update_instance_metadata_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + gce_setup=notebooks_v2.GceSetup( + metadata={ + "idle-timeout-seconds": "3600", + "version": "20261101-2130", + "notebooks-api-version": "v2", + } + ), + ) + mock_client.update_instance.return_value = operation + + update_instance_metadata.update_instance_metadata( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + request = mock_client.update_instance.call_args.kwargs["request"] + assert request.instance.name == INSTANCE_NAME + metadata = dict(request.instance.gce_setup.metadata) + assert metadata == {"idle-timeout-seconds": "3600"} + assert list(request.update_mask.paths) == ["gce_setup.metadata"] + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Idle Timeout Seconds: 3600" in captured.out + assert "Image Version: 20261101-2130" in captured.out + assert "Metadata Key Count: 3" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_update_instance_metadata_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.update_instance.side_effect = exceptions.NotFound("Instance not found") + + update_instance_metadata.update_instance_metadata( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/update_instance_shielded_config_test.py b/notebooks/v2/tests/update_instance_shielded_config_test.py new file mode 100644 index 00000000000..5818955aba3 --- /dev/null +++ b/notebooks/v2/tests/update_instance_shielded_config_test.py @@ -0,0 +1,91 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import update_instance_shielded_config + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_update_instance_shielded_config_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.STOPPED, + gce_setup=notebooks_v2.GceSetup( + shielded_instance_config=notebooks_v2.ShieldedInstanceConfig( + enable_secure_boot=True, + enable_vtpm=True, + enable_integrity_monitoring=True, + ) + ), + ) + mock_client.update_instance.return_value = operation + + update_instance_shielded_config.update_instance_shielded_config( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + request = mock_client.update_instance.call_args.kwargs["request"] + assert request.instance.name == INSTANCE_NAME + shielded = request.instance.gce_setup.shielded_instance_config + assert shielded.enable_secure_boot + assert shielded.enable_vtpm + assert shielded.enable_integrity_monitoring + assert sorted(request.update_mask.paths) == sorted( + [ + "gce_setup.shielded_instance_config.enable_secure_boot", + "gce_setup.shielded_instance_config.enable_vtpm", + "gce_setup.shielded_instance_config.enable_integrity_monitoring", + ] + ) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Secure Boot Enabled: True" in captured.out + assert "State: STOPPED" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_update_instance_shielded_config_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.update_instance.side_effect = exceptions.NotFound("Instance not found") + + update_instance_shielded_config.update_instance_shielded_config( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/upgrade_instance_test.py b/notebooks/v2/tests/upgrade_instance_test.py new file mode 100644 index 00000000000..58aae0f3748 --- /dev/null +++ b/notebooks/v2/tests/upgrade_instance_test.py @@ -0,0 +1,68 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +from .. import upgrade_instance + +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_upgrade_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + ) + mock_client.upgrade_instance.return_value = operation + + upgrade_instance.upgrade_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + request = mock_client.upgrade_instance.call_args.kwargs["request"] + assert request.name == INSTANCE_NAME + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) +def test_upgrade_instance_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.upgrade_instance.side_effect = exceptions.NotFound("Instance not found") + + upgrade_instance.upgrade_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/update_instance_machine_type.py b/notebooks/v2/update_instance_machine_type.py new file mode 100644 index 00000000000..b481ae616ff --- /dev/null +++ b/notebooks/v2/update_instance_machine_type.py @@ -0,0 +1,62 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_update_instance_machine_type] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +from google.protobuf import field_mask_pb2 + + +def update_instance_machine_type( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", + machine_type: str = "n1-standard-8", +) -> None: + """Updates the machine type of a stopped Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the stopped Workbench instance to update. + machine_type: The target Compute Engine machine type. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + gce_setup = notebooks_v2.GceSetup(machine_type=machine_type) + instance = notebooks_v2.Instance(name=name, gce_setup=gce_setup) + update_mask = field_mask_pb2.FieldMask(paths=["gce_setup.machine_type"]) + request = notebooks_v2.UpdateInstanceRequest( + instance=instance, + update_mask=update_mask, + ) + + try: + operation = client.update_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result(timeout=300) + print(f"Resource ID: {response.name}") + print(f"Machine Type: {response.gce_setup.machine_type}") + print(f"State: {response.state.name}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_update_instance_machine_type] diff --git a/notebooks/v2/update_instance_metadata.py b/notebooks/v2/update_instance_metadata.py new file mode 100644 index 00000000000..ae2f01e9c08 --- /dev/null +++ b/notebooks/v2/update_instance_metadata.py @@ -0,0 +1,66 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_update_instance_metadata] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +from google.protobuf import field_mask_pb2 + + +def update_instance_metadata( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Updates custom metadata on a Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to update. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + gce_setup = notebooks_v2.GceSetup( + metadata={"idle-timeout-seconds": "3600"}, + ) + instance = notebooks_v2.Instance(name=name, gce_setup=gce_setup) + update_mask = field_mask_pb2.FieldMask(paths=["gce_setup.metadata"]) + request = notebooks_v2.UpdateInstanceRequest( + instance=instance, + update_mask=update_mask, + ) + + try: + operation = client.update_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result(timeout=300) + # The service merges the update into metadata it manages itself, such + # as the image version, so only the updated key changes. + metadata = response.gce_setup.metadata + print(f"Resource ID: {response.name}") + print(f"Idle Timeout Seconds: {metadata.get('idle-timeout-seconds')}") + print(f"Image Version: {metadata.get('version')}") + print(f"Metadata Key Count: {len(metadata)}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_update_instance_metadata] diff --git a/notebooks/v2/update_instance_shielded_config.py b/notebooks/v2/update_instance_shielded_config.py new file mode 100644 index 00000000000..a8d0cbc349c --- /dev/null +++ b/notebooks/v2/update_instance_shielded_config.py @@ -0,0 +1,72 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_update_instance_shielded_config] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +from google.protobuf import field_mask_pb2 + + +def update_instance_shielded_config( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Updates the Shielded VM configuration on a stopped Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the stopped Workbench instance to update. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + shielded_config = notebooks_v2.ShieldedInstanceConfig( + enable_secure_boot=True, + enable_vtpm=True, + enable_integrity_monitoring=True, + ) + gce_setup = notebooks_v2.GceSetup(shielded_instance_config=shielded_config) + instance = notebooks_v2.Instance(name=name, gce_setup=gce_setup) + update_mask = field_mask_pb2.FieldMask( + paths=[ + "gce_setup.shielded_instance_config.enable_secure_boot", + "gce_setup.shielded_instance_config.enable_vtpm", + "gce_setup.shielded_instance_config.enable_integrity_monitoring", + ] + ) + request = notebooks_v2.UpdateInstanceRequest( + instance=instance, + update_mask=update_mask, + ) + + try: + operation = client.update_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result(timeout=300) + secure_boot = response.gce_setup.shielded_instance_config.enable_secure_boot + print(f"Resource ID: {response.name}") + print(f"Secure Boot Enabled: {secure_boot}") + print(f"State: {response.state.name}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_update_instance_shielded_config] diff --git a/notebooks/v2/upgrade_instance.py b/notebooks/v2/upgrade_instance.py new file mode 100644 index 00000000000..654d03cf753 --- /dev/null +++ b/notebooks/v2/upgrade_instance.py @@ -0,0 +1,52 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START notebooks_v2_upgrade_instance] +import sys + +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def upgrade_instance( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Upgrades a Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to upgrade. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + request = notebooks_v2.UpgradeInstanceRequest(name=name) + + try: + operation = client.upgrade_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result(timeout=600) + print(f"Resource ID: {response.name}") + print(f"State: {response.state.name}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_upgrade_instance]