sagemaker module

Author: luigift

awswrangler/sagemaker.py

@@ -0,0 +1,39 @@
+import pickle
+import tarfile
+import logging
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class SageMaker:
+    def __init__(self, session):
+        self._session = session
+        self._client_s3 = session.boto3_session.client(service_name="s3", use_ssl=True, config=session.botocore_config)
+
+    @staticmethod
+    def _parse_path(path):
+        path2 = path.replace("s3://", "")
+        parts = path2.partition("/")
+        return parts[0], parts[2]
+
+    def get_job_outputs(self, path: str) -> Any:
+
+        bucket, key = SageMaker._parse_path(path)
+        if key.split("/")[-1] != "model.tar.gz":
+            key = f"{key}/model.tar.gz"
+        body = self._client_s3.get_object(Bucket=bucket, Key=key)["Body"].read()
+        body = tarfile.io.BytesIO(body)
+        tar = tarfile.open(fileobj=body)
+
+        results = []
+        for member in tar.getmembers():
+            f = tar.extractfile(member)
+            file_type = member.name.split(".")[-1]
+
+            if file_type == "pkl":
+                f = pickle.load(f)
+
+            results.append(f)
+
+        return results

awswrangler/session.py

@@ -13,6 +13,8 @@
 from awswrangler.glue import Glue
 from awswrangler.redshift import Redshift
 from awswrangler.emr import EMR
+from awswrangler.sagemaker import SageMaker
+
 
 PYSPARK_INSTALLED = False
 if importlib.util.find_spec("pyspark"):  # type: ignore
@@ -112,6 +114,7 @@ def __init__(self,
         self._glue = None
         self._redshift = None
         self._spark = None
+        self._sagemaker = None
 
     def _load_new_boto3_session(self):
         """
@@ -281,6 +284,12 @@ def redshift(self):
             self._redshift = Redshift(session=self)
         return self._redshift
 
+    @property
+    def sagemaker(self):
+        if not self._sagemaker:
+            self._sagemaker = SageMaker(session=self)
+        return self._sagemaker
+
     @property
     def spark(self):
         if not PYSPARK_INSTALLED:

requirements-dev.txt

@@ -3,8 +3,10 @@ mypy~=0.750
 flake8~=3.7.9
 pytest-cov~=2.8.1
 cfn-lint~=0.26.0
+scikit-learn==0.22
+sklearn==0.0
 twine~=3.1.1
 wheel~=0.33.6
 sphinx~=2.2.2
 pyspark~=2.4.4
-pyspark-stubs~=2.4.0.post6
+pyspark-stubs~=2.4.0.post6

testing/test_awswrangler/test_sagemaker.py

@@ -0,0 +1,58 @@
+import os
+import pickle
+import logging
+import tarfile
+
+import boto3
+import pytest
+
+from awswrangler import Session
+from sklearn.linear_model import LinearRegression
+
+logging.basicConfig(level=logging.INFO, format="[%(asctime)s][%(levelname)s][%(name)s][%(funcName)s] %(message)s")
+logging.getLogger("awswrangler").setLevel(logging.DEBUG)
+
+
+@pytest.fixture(scope="module")
+def session():
+    yield Session()
+
+
+@pytest.fixture(scope="module")
+def cloudformation_outputs():
+    response = boto3.client("cloudformation").describe_stacks(StackName="aws-data-wrangler-test-arena")
+    outputs = {}
+    for output in response.get("Stacks")[0].get("Outputs"):
+        outputs[output.get("OutputKey")] = output.get("OutputValue")
+    yield outputs
+
+
+@pytest.fixture(scope="module")
+def bucket(session, cloudformation_outputs):
+    if "BucketName" in cloudformation_outputs:
+        bucket = cloudformation_outputs["BucketName"]
+        session.s3.delete_objects(path=f"s3://{bucket}/")
+    else:
+        raise Exception("You must deploy the test infrastructure using Cloudformation!")
+    yield bucket
+    session.s3.delete_objects(path=f"s3://{bucket}/")
+
+
+def test_get_job_outputs(session, bucket):
+    model_path = "output"
+    s3 = boto3.resource("s3")
+
+    lr = LinearRegression()
+    with open("model.pkl", "wb") as fp:
+        pickle.dump(lr, fp, pickle.HIGHEST_PROTOCOL)
+
+    with tarfile.open("model.tar.gz", "w:gz") as tar:
+        tar.add("model.pkl")
+
+    s3.Bucket(bucket).upload_file("model.tar.gz", f"{model_path}/model.tar.gz")
+    outputs = session.sagemaker.get_job_outputs(f"{bucket}/{model_path}")
+
+    os.remove("model.pkl")
+    os.remove("model.tar.gz")
+
+    assert type(outputs[0]) == LinearRegression