Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Patching for generativeai Python package #1329

Merged
merged 22 commits into from
Dec 3, 2023
Merged

Patching for generativeai Python package #1329

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
bb9b855
Patching for Python package
psbang Nov 29, 2023
6508445
Typo fix
psbang Nov 29, 2023
684dc6a
Run python script instead of just adding it to docker image
psbang Nov 30, 2023
a49c250
Try using ADD and CMD
psbang Nov 30, 2023
de9912c
Add logic to sitecustomize
psbang Nov 30, 2023
3477d4d
Escape path
psbang Nov 30, 2023
a3c4b16
Address comments
psbang Nov 30, 2023
f301127
Merge branch 'main' into psbang-genaipatch
psbang Nov 30, 2023
4c4240a
Check for another env variable
psbang Nov 30, 2023
5e53b3d
Merge branch 'main' into psbang-genaipatch
psbang Dec 1, 2023
5a1504a
Address comments
psbang Dec 1, 2023
4c5f7cd
Build fix hopefully
psbang Dec 1, 2023
7d4f55d
Add unit test for headers
psbang Dec 1, 2023
a06ec16
bug fix
psbang Dec 1, 2023
0f7eaa5
Fix import issue and test
psbang Dec 1, 2023
f1559e0
small modification
psbang Dec 1, 2023
f89b6a2
add a test for disabled proxy for generativeai api
Philmod Dec 2, 2023
a9d45f9
remove print
Philmod Dec 2, 2023
f7b69ea
change name of test
Philmod Dec 2, 2023
dfb9145
remove disabled test to check if that's the reason it's broken on ci
Philmod Dec 3, 2023
ff335d6
readd disabled test
Philmod Dec 3, 2023
10b9a81
readd disabled test
Philmod Dec 3, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions Dockerfile.tmpl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -634,6 +634,10 @@ RUN mkdir -p ~/src && git clone https://github.com/SohierDane/BigQuery_Helper ~/
sed -i 's/)/packages=["bq_helper"])/g' ~/src/BigQuery_Helper/setup.py && \
pip install -e ~/src/BigQuery_Helper && \
/tmp/clean-layer.sh

RUN pip install wrapt \
google-generativeai && \
/tmp/clean-layer.sh

# Add BigQuery client proxy settings
ENV PYTHONUSERBASE "/root/.local"
Expand All @@ -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"]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

e.g.

docker-python/patches/sitecustomize.py

Lines 74 to 75 in 823af9a

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


# Add Kaggle module resolver
ADD patches/kaggle_module_resolver.py /opt/conda/lib/python3.10/site-packages/tensorflow_hub/kaggle_module_resolver.py
RUN sed -i '/from tensorflow_hub import uncompressed_module_resolver/a from tensorflow_hub import kaggle_module_resolver' /opt/conda/lib/python3.10/site-packages/tensorflow_hub/config.py && \
Expand Down
24 changes: 24 additions & 0 deletions patches/generativeaipatch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
import wrapt
import os

@wrapt.when_imported('google.generativeai')
def post_import_logic(module):
old_configure = module.configure
Philmod marked this conversation as resolved.
Show resolved Hide resolved
def new_configure(*args, **kwargs):
if ('default_metadata' in kwargs):
default_metadata = kwargs['default_metadata']
else:
default_metadata = []
kwargs['transport'] = 'rest' # Only support REST requests for now
default_metadata.append(("x-kaggle-proxy-data", os.environ['KAGGLE_DATA_PROXY_TOKEN']))
Philmod marked this conversation as resolved.
Show resolved Hide resolved
default_metadata.append(('x-kaggle-authorization', f'Bearer {os.environ['KAGGLE_USER_SECRETS_TOKEN']}'))
kwargs['default_metadata'] = default_metadata
if ('client_options' in kwargs):
client_options = kwargs['client_options']
else:
client_options = {}
client_options['api_endpoint'] = os.environ['KAGGLE_DATA_PROXY_URL'] + '/palmapi'
Philmod marked this conversation as resolved.
Show resolved Hide resolved
kwargs['client_options'] = client_options
old_configure(*args, **kwargs)
module.configure = new_configure
module.configure() # generativeai can use GOOGLE_API_KEY env variable, so make sure we have the other configs set