Skip to content

Refactor sagemaker.session.Session #1463

Closed
@laurenyu

Description

@laurenyu

Is your feature request related to a problem? Please describe.
We’ve thus far treated Session as a class for handling all AWS calls, which has led to a class that’s doing far too much.

Describe the solution you'd like
I'm proposing a complete overhaul of the class. The Session class’s sole purpose is to keep defaults and user settings for using the Python SDK. This includes the boto3 session, as well as the Local Mode config and default bucket. I think there’s potential here to extend this to other user settings, e.g. input/output S3 prefixes, instance types.

Similar to how boto3’s Session allows the user to create clients for various AWS clients and resources, our new Session will lead to clients for calling SageMaker’s various services, as well as S3. This means all of the S3 functions go in one client, all training job functions go in another, all tuning job functions go in a third, etc.

I’m also proposing that these "SageMaker clients" be private for two reasons. First, the interface we want to be used is the one defined by our estimators, predictors, etc. These classes are how we interact with SageMaker, and we shouldn’t be asking users to have to instantiate too many objects just to, say, kick off a job.

The second reason is that we then have flexibility to change our functions as the SageMaker API evolves. With multi-algo tuning, we had to create a second function for creating tuning jobs that was almost identical to the first because changing the first one would have been a breaking change.

Lastly, there are two additional classes defined in session.py that are not used in that file. These will be moved:

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions