Skip to content

build(setup): pin sagemaker<3 so import cohere_aws works again - #234

Open
tsushanth wants to merge 1 commit into
cohere-ai:mainfrom
tsushanth:fix/issue-233-pin-sagemaker-lt-3
Open

build(setup): pin sagemaker<3 so import cohere_aws works again#234
tsushanth wants to merge 1 commit into
cohere-ai:mainfrom
tsushanth:fix/issue-233-pin-sagemaker-lt-3

Conversation

@tsushanth

Copy link
Copy Markdown

Summary

Closes #233.

sagemaker 3.0 removed the sagemaker.s3 module — S3Downloader, S3Uploader, and parse_s3_url are no longer importable from that path. cohere_aws/client.py:12 imports all three at module load:

from sagemaker.s3 import S3Downloader, S3Uploader, parse_s3_url

So any environment that resolves to sagemaker>=3 throws ModuleNotFoundError on the first import cohere_aws, long before the caller has a chance to do anything else:

>>> import cohere_aws
ModuleNotFoundError: No module named 'sagemaker.s3'

setup.py declared the dep as unpinned sagemaker, which is now wrong on any fresh install that picks the latest published version.

Fix

-                 install_requires=['boto3>=1.34.162', 'sagemaker'],
+                 # sagemaker 3.0 removed the `sagemaker.s3` module
+                 # (S3Downloader / S3Uploader / parse_s3_url), which client.py
+                 # imports unconditionally — installing this package alongside
+                 # sagemaker>=3 raises ModuleNotFoundError on first import.
+                 # Pin to <3 until the imports are migrated to sagemaker v3's
+                 # surface. See https://github.com/cohere-ai/cohere-aws/issues/233.
+                 install_requires=['boto3>=1.34.162', 'sagemaker<3'],

Matches what @robertkerr-cohere suggested on the issue thread:

This is caused by using sagemaker version 3, which removed the sagemaker.s3 path. This package should pin sagemaker<3 or upgrade code for v3.

This PR takes the first half of his suggestion — restore importability today. The longer-term migration to sagemaker v3's APIs (where S3Downloader / S3Uploader moved to sagemaker.s3_utils, etc.) is a separate, larger change worth tracking in its own issue.

The inline comment links #233 so a future maintainer trying to lift the constraint sees the migration prerequisite right there.

Test plan

  • Manual: a fresh pip install cohere-aws against this branch now resolves sagemaker<3 and import cohere_aws succeeds. On main the same flow either fails at install (if pip happened to pick sagemaker>=3) or at first import (if a newer sagemaker was already in the env).
  • boto3>=1.34.162 constraint left untouched — only the sagemaker dep is changed.

Closes cohere-ai#233.

sagemaker 3.0 removed the `sagemaker.s3` module, taking `S3Downloader`,
`S3Uploader`, and `parse_s3_url` with it. `cohere_aws/client.py:12`
imports all three at module load:

    from sagemaker.s3 import S3Downloader, S3Uploader, parse_s3_url

so any environment that resolves to sagemaker>=3 throws
ModuleNotFoundError on the first `import cohere_aws` — long before the
caller has a chance to do anything else.

setup.py declared the dep as unpinned `sagemaker`, which is now wrong
on any fresh install that picks the latest published version. Pin to
<3 so installs continue to work; an inline comment links the issue and
states that the long-term fix is to migrate the imports to sagemaker's
v3 surface, at which point the constraint can be lifted.

Matches what @robertkerr-cohere suggested on the issue thread.
@tsushanth
tsushanth requested a review from a team as a code owner June 12, 2026 15:12
@tsushanth
tsushanth requested a review from jpekmez June 12, 2026 15:12
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Import cohere_ws giving an issue No module named 'sagemaker.s3'

1 participant