def get_session()

in repos/serving/infra/sm_pipeline_utils.py [0:0]

• 20 lines of code
• 1 McCabe index (conditional complexity)

def get_session(region: str, default_bucket: str) -> sagemaker.session.Session:
    """Gets the sagemaker session based on the region.
    Args:
        region: the aws region to start the session
        default_bucket: the bucket to use for storing the artifacts

    Returns:
        `sagemaker.session.Session instance
    """

    boto_session = boto3.Session(region_name=region)

    sagemaker_client = boto_session.client("sagemaker")
    runtime_client = boto_session.client("sagemaker-runtime")
    try:
        sagemaker.session.Session(
            boto_session=boto_session,
            sagemaker_client=sagemaker_client,
            sagemaker_runtime_client=runtime_client,
            default_bucket=default_bucket,
        )
        logger.info("SageMaker Session created")
    except:
        logger.exception("Failed to generate a SageMaker Session")

    return sagemaker.session.Session(
        boto_session=boto_session,
        sagemaker_client=sagemaker_client,
        sagemaker_runtime_client=runtime_client,
        default_bucket=default_bucket,
    )