Patching for generativeai Python package #1329

psbang · 2023-11-29T21:19:45Z

A patch for the generativeai Python package that makes sure all requests are directed through our data proxy and configurations are set accordingly.

http://b/308644984

Philmod

Please write unit tests to make sure the new package and your patch work

patches/generativeaipatch.py

rosbo · 2023-11-30T19:33:41Z

Dockerfile.tmpl

@@ -647,6 +651,10 @@ ADD patches/sitecustomize.py /root/.local/lib/python3.10/site-packages/sitecusto
 # Override default imagemagick policies
 ADD patches/imagemagick-policy.xml /etc/ImageMagick-6/policy.xml

+# Add generativeai Python patch
+ADD patches/generativeaipatch.py /root/.local/lib/python3.10/site-packages/generativeaipatch.py
+CMD ["python", "/root/.local/lib/python3.10/site-packages/generativeaipatch.py"]


CMD is the Dockerfile instruction to update the command that is run when you start a container with the image. Do not update the CMD.

You can add logic that needs to be called to sitecustomize.py which gets call when a notebook session starts up.

e.g.

docker-python/patches/sitecustomize.py

Lines 74 to 75 in 823af9a

if not hasattr(sys, 'frozen'):

sys.meta_path.insert(0, GcpModuleFinder())

patches/generativeaipatch.py

rosbo

Add a test that calls a fake HTTP client and assert that the headers are set properly.

Example test with a test HTTP server: https://github.com/Kaggle/kagglehub/blob/66f2121adccae70f05c72594113a7e463fa6ddce/tests/test_http_model_download.py#L19

rosbo · 2023-11-30T20:03:37Z

patches/sitecustomize.py

@@ -73,3 +75,27 @@ def exec_module(self, module):

 if not hasattr(sys, 'frozen'):
    sys.meta_path.insert(0, GcpModuleFinder())
+
+@wrapt.when_imported('google.generativeai')


Is this called only the first time this module is imported or again every time?

To confirm, you can simply add a print statement and test in a notebook by running "import 'google.generativeai" a few times.

It's only called the first time it seems.

rosbo · 2023-11-30T21:08:46Z

Before merging, make sure to also test in the kernels editor using the psbang-genaipatch-pretest image tag like we discussed in our 1:1 earlier this week :)

patches/sitecustomize.py

Philmod · 2023-12-01T09:15:36Z

patches/sitecustomize.py

+            default_metadata = kwargs['default_metadata']
+        else:
+            default_metadata = []
+        kwargs['transport'] = 'rest' # Only support REST requests for now


Because this is hard coded here, how are you going to enable grpc when it's ready? All the users using this image will have this code forever.

I see, I'll set an environment variable then.

Philmod · 2023-12-01T09:16:51Z

patches/sitecustomize.py

+            client_options = kwargs['client_options']
+        else:
+            client_options = {}
+        client_options['api_endpoint'] = os.environ['KAGGLE_DATA_PROXY_URL'] + '/palmapi'


Same: hard coding this endpoint here means you will never be able to change it without breaking all the users using this image

patches/sitecustomize.py

Philmod · 2023-12-01T16:14:45Z

Please also confirm that installing another version of the package doesn't break your import patch:

!pip install ... ==another-version

import ...
# patch still works

Philmod · 2023-12-02T07:24:29Z

Please write unit tests to make sure the new package and your patch work

I'm guessing you were able to test this image and it worked as expected?

Philmod · 2023-12-03T18:37:25Z

@psbang I'm merging this PR for you to get ready for a release tomorrow morning.

psbang added 3 commits November 29, 2023 20:42

Patching for Python package

bb9b855

Typo fix

6508445

Run python script instead of just adding it to docker image

684dc6a

Philmod suggested changes Nov 30, 2023

View reviewed changes

patches/generativeaipatch.py Outdated Show resolved Hide resolved

Try using ADD and CMD

a49c250

psbang requested a review from rosbo November 30, 2023 19:20

rosbo requested changes Nov 30, 2023

View reviewed changes

rosbo reviewed Nov 30, 2023

View reviewed changes

patches/generativeaipatch.py Outdated Show resolved Hide resolved

psbang added 2 commits November 30, 2023 19:59

Add logic to sitecustomize

de9912c

Escape path

3477d4d

rosbo reviewed Nov 30, 2023

View reviewed changes

patches/generativeaipatch.py Outdated Show resolved Hide resolved

rosbo reviewed Nov 30, 2023

View reviewed changes

psbang added 4 commits November 30, 2023 22:31

Address comments

a3c4b16

Merge branch 'main' into psbang-genaipatch

f301127

Check for another env variable

4c4240a

Merge branch 'main' into psbang-genaipatch

5e53b3d

Philmod reviewed Dec 1, 2023

View reviewed changes

patches/sitecustomize.py Outdated Show resolved Hide resolved

Philmod reviewed Dec 1, 2023

View reviewed changes

patches/sitecustomize.py Show resolved Hide resolved

Address comments

5a1504a

psbang added 5 commits December 1, 2023 19:18

Build fix hopefully

4c5f7cd

Add unit test for headers

7d4f55d

bug fix

a06ec16

Fix import issue and test

0f7eaa5

small modification

f1559e0

add a test for disabled proxy for generativeai api

f89b6a2

Philmod approved these changes Dec 2, 2023

View reviewed changes

Philmod added 5 commits December 2, 2023 14:19

remove print

a9d45f9

change name of test

f7b69ea

remove disabled test to check if that's the reason it's broken on ci

dfb9145

readd disabled test

ff335d6

readd disabled test

10b9a81

Philmod merged commit 3fdce73 into main Dec 3, 2023
3 checks passed

Philmod deleted the psbang-genaipatch branch December 3, 2023 18:40

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Patching for generativeai Python package #1329

Patching for generativeai Python package #1329

psbang commented Nov 29, 2023 •

edited by Philmod

Loading

Philmod left a comment

rosbo Nov 30, 2023

rosbo Nov 30, 2023

rosbo left a comment

rosbo Nov 30, 2023 •

edited

Loading

psbang Nov 30, 2023

rosbo commented Nov 30, 2023

Philmod Dec 1, 2023

psbang Dec 1, 2023

Philmod Dec 1, 2023

Philmod commented Dec 1, 2023 •

edited

Loading

Philmod commented Dec 2, 2023

Philmod commented Dec 3, 2023

	if not hasattr(sys, 'frozen'):
	sys.meta_path.insert(0, GcpModuleFinder())

Patching for generativeai Python package #1329

Patching for generativeai Python package #1329

Conversation

psbang commented Nov 29, 2023 • edited by Philmod Loading

Philmod left a comment

Choose a reason for hiding this comment

rosbo Nov 30, 2023

Choose a reason for hiding this comment

rosbo Nov 30, 2023

Choose a reason for hiding this comment

rosbo left a comment

Choose a reason for hiding this comment

rosbo Nov 30, 2023 • edited Loading

Choose a reason for hiding this comment

psbang Nov 30, 2023

Choose a reason for hiding this comment

rosbo commented Nov 30, 2023

Philmod Dec 1, 2023

Choose a reason for hiding this comment

psbang Dec 1, 2023

Choose a reason for hiding this comment

Philmod Dec 1, 2023

Choose a reason for hiding this comment

Philmod commented Dec 1, 2023 • edited Loading

Philmod commented Dec 2, 2023

Philmod commented Dec 3, 2023

psbang commented Nov 29, 2023 •

edited by Philmod

Loading

rosbo Nov 30, 2023 •

edited

Loading

Philmod commented Dec 1, 2023 •

edited

Loading