sagemaker increases import time from ~1 sec to ~10 seconds

[ben-albrecht]

SDK Version(s) & bug reproducer

import time
t = time.time()
import cohere
print(time.time() - t)

print(cohere.__version__)

10.188275575637817
5.11.0

0.7559852600097656
5.10.0

This also significantly impacts import time of packages that depend on cohere, such as instructor-ai/instructor.

[neilkumar]

It's also adding a lot of noise to the logs

[10/16/24 02:15:29] INFO     INFO  Not applying SDK defaults from location: /etc/xdg/sagemaker/config.yaml                                                                                                                                                                            config_utils.py:243
                    INFO     INFO  Not applying SDK defaults from location: /home/app/.config/sagemaker/config.yaml       

[andaldanar]

Facing the same issue. Just want to use the Rerank models directly from the Cohere API, but the sagemaker import (and others like mlflow, docker, etc.) are dealbreakers for our use case.

It'd be great if we could install only the core Cohere platform stuff by default, with optional installs for specific platforms (Bedrock, SageMaker, Azure, GCP, Oracle OCI)

This way, we can avoid the slow imports and unnecessary dependencies when we're just using the basic Cohere API. Could we have a lite/ core version that doesn't pull in all these extra packages?

Thank you!

[ben-albrecht]

To speed up imports until this is fixed on the next release, you can switch to previous release:

uv pip install cohere==5.10.0

[lucasgadams]

+1, it is definitely a pain now with the sagemaker dependency. If you set logging to debug you can see it is making a ton of network requests at import time? I guess it is trying to get some configuration stuff and is assuming you are importing sagemaker on an aws machine. I see logs like this being generated:

2024-10-18T14:16:01.962101Z [debug    ] Caught retryable HTTP exception while making metadata service request to http://169.254.169.254/latest/api/token: Could not connect to the endpoint URL: "http://169.254.169.254/latest/api/token" [botocore.utils]
botocore.exceptions.EndpointConnectionError: Could not connect to the endpoint URL: "http://169.254.169.254/latest/meta-data/iam/security-credentials/"
...

with a few retries. To be clear, this happens when you import sagemaker alone so doesn't seem an issue with cohere besides the dependency. Not sure why the sagemaker library is designed like that.

[ben-albrecht]

@billytrend-cohere @1vn - I believe the issue began with #588, with the introduction of sagemaker as a dependency.

Perhaps you could make this an optional dependency, such that the default install does not increase import time by ~10 seconds?

e.g.

uv pip install cohere[sagemaker]

[billytrend-cohere]

Hey @ben-albrecht thanks for flagging this. I've made the aws-related dependencies optional now as of 5.11.1. I will try out your script to confirm that the issue is resolved! Thanks

[billytrend-cohere]

Ok I think we're good but please re-open if you have related issues!

I see:

1.8758995532989502
5.11.1

sagemaker.config INFO - Not applying SDK defaults from location: /etc/xdg/sagemaker/config.yaml
sagemaker.config INFO - Not applying SDK defaults from location: /home/codespace/.config/sagemaker/config.yaml
2.3514339923858643
5.11.0