iterative · daavoo · Aug 16, 2023 · Aug 10, 2023 · Aug 10, 2023 · Aug 10, 2023
diff --git a/.devcontainer.json b/.devcontainer.json
@@ -1,6 +1,7 @@
 {
   "name": "example-repos-dev",
   "image": "mcr.microsoft.com/devcontainers/python:3.10",
+  "runArgs": ["--ipc=host"],
   "extensions": ["Iterative.dvc", "ms-python.python", "redhat.vscode-yaml"],
   "features": {
       "ghcr.io/devcontainers/features/nvidia-cuda:1": {

diff --git a/example-get-started-experiments/code/.github/workflows/deploy-model.yml b/example-get-started-experiments/code/.github/workflows/deploy-model.yml
@@ -0,0 +1,70 @@
+name: Deploy model
+
+on:
+  push:
+    # When a new version is registered in Studio Model Registry
+    tags:
+      - "results/train=pool-segmentation#*"
+
+  workflow_dispatch:
+    inputs:
+      version:
+        description: 'Manual version name'
+        required: true
+        type: string
+
+permissions:
+  contents: write
+  id-token: write
+
+jobs:
+  parse:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        fetch-depth: 0
+    - name: "Parse GTO tag"
+      id: gto
+      uses: iterative/gto-action@14723404a00bb0c1e759c02ffcd24279df5815c2
+    outputs:
+      event: ${{ steps.gto.outputs.event }}
+      name: ${{ steps.gto.outputs.name }}
+      stage: ${{ steps.gto.outputs.stage }}
+      version: ${{ steps.gto.outputs.version }}
+
+  deploy-model:
+    needs: parse
+    if: ${{ needs.parse.outputs.event }} == 'assignment'
+    environment: cloud
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        fetch-depth: 0
+
+    - uses: aws-actions/configure-aws-credentials@v2
+      with:
+        aws-region: us-east-2
+        role-to-assume: ${{ secrets.AWS_ROLE_TO_ASSUME }}
+        role-duration-seconds: 43200
+
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: '3.8'
+        cache: 'pip'
+        cache-dependency-path: requirements.txt
+
+    - run: pip install -r requirements.txt
+
+    - run:  dvc remote add -d --local storage s3://dvc-public/remote/get-started-pools
+
+    - run: |
+        MODEL_DATA=$(dvc get --show-url . model.tar.gz)
+        python sagemaker/deploy_model.py \
+        --name ${{ needs.parse.outputs.name }} \
+        --stage ${{ needs.parse.outputs.stage }} \
+        --version ${{ needs.parse.outputs.version }} \
+        --model_data $MODEL_DATA \
+        --role ${{ secrets.AWS_ROLE_TO_ASSUME }}
diff --git a/example-get-started-experiments/code/.github/workflows/dvc-studio.yml b/example-get-started-experiments/code/.github/workflows/dvc-studio.yml
@@ -62,7 +62,7 @@ jobs:
       - uses: aws-actions/configure-aws-credentials@v1
         with:
           aws-region: us-east-2
-          role-to-assume: arn:aws:iam::342840881361:role/SandboxUser
+          role-to-assume: ${{ secrets.AWS_ROLE_TO_ASSUME }}
           role-duration-seconds: 43200
       - name: Create Runner
         env:
@@ -90,7 +90,7 @@ jobs:
       - uses: aws-actions/configure-aws-credentials@v1
         with:
           aws-region: us-east-2
-          role-to-assume: arn:aws:iam::342840881361:role/SandboxUser
+          role-to-assume: ${{ secrets.AWS_ROLE_TO_ASSUME }}
           role-duration-seconds: 43200
 
       - run: pip install -r requirements.txt

diff --git a/example-get-started-experiments/code/README.md b/example-get-started-experiments/code/README.md
@@ -108,16 +108,11 @@ This tag also contains a GitHub Actions workflow that reruns the pipeline if any
 [CML](https://cml.dev/) is used in this workflow to provision a cloud-based GPU 
 machine as well as report model performance results in Pull Requests.
 
-## Deploying the model
-
-Check out the [PR](https://github.com/iterative/example-get-started-experiments/pulls)
-that adds this model to 
-[Iterative Studio Model Registry](https://dvc.org/doc/studio/user-guide/model-registry/what-is-a-model-registry).
-You can [trigger CI/CD](https://dvc.org/doc/studio/user-guide/model-registry/use-models#deploying-and-publishing-models-in-cicd)
-by [registering versions](https://dvc.org/doc/studio/user-guide/model-registry/register-version)
-and [assigning stages](https://dvc.org/doc/studio/user-guide/model-registry/assign-stage)
-in Model Registry, building and publishing Docker images with the model,
-or deploying the model to the cloud.
+## Model Deployment
+
+Check out the [GitHub Workflow](https://github.com/iterative/example-get-started-experiments/blob/main/.github/workflows/deploy-model.yml)
+that uses the [Iterative Studio Model Registry](https://dvc.org/doc/studio/user-guide/model-registry/what-is-a-model-registry).
+to deploy the model to [AWS Sagemaker](https://aws.amazon.com/es/sagemaker/) whenever a new [version is registered](https://dvc.org/doc/studio/user-guide/model-registry/register-version).
 
 ## Project structure
 

diff --git a/example-get-started-experiments/code/requirements.txt b/example-get-started-experiments/code/requirements.txt
@@ -1,4 +1,5 @@
 dvc[s3]>=3.0
 dvclive>=2.11.3
 fastai
-python-box
+python-box
+sagemaker
diff --git a/example-get-started-experiments/code/sagemaker/code/inference.py b/example-get-started-experiments/code/sagemaker/code/inference.py
@@ -0,0 +1,50 @@
+"""
+Reference:
+https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#id4
+"""
+import io
+import os
+
+import numpy as np
+import torch
+from PIL import Image
+from torchvision.transforms import Compose, Normalize, Resize, ToTensor
+
+
+def model_fn(model_dir, context):
+    kwargs = {
+        "f": os.path.join(model_dir, "code/model.pth")
+    }
+    if not torch.cuda.is_available():
+        kwargs["map_location"] = torch.device("cpu")
+    model = torch.load(**kwargs)
+    return model
+
+
+def input_fn(request_body, request_content_type, context):
+    if request_content_type:
+        img_pil = Image.open(io.BytesIO(request_body))
+        img_transform = Compose([Resize(512), ToTensor(), Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
+        img_tensor = img_transform(img_pil).unsqueeze_(0)
+        return img_tensor
+    else:
+        raise ValueError(f"Unsupported request_content_type {request_content_type}")
+
+
+def predict_fn(input_object, model, context):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    with torch.no_grad():
+        result = model(input_object)
+    return result
+
+
+def output_fn(prediction_output, content_type):
+    output = np.array(
+        prediction_output[:, 1, :] > 0.5, dtype=np.uint8
+    )
+    if torch.cuda.is_available():
+        output = output.cpu()
+    buffer = io.BytesIO()
+    np.save(buffer, output)
+    return buffer.getvalue()
diff --git a/example-get-started-experiments/code/sagemaker/code/requirements.txt b/example-get-started-experiments/code/sagemaker/code/requirements.txt
@@ -0,0 +1,4 @@
+fastai
+pillow
+torch
+torchvision
diff --git a/example-get-started-experiments/code/sagemaker/deploy_model.py b/example-get-started-experiments/code/sagemaker/deploy_model.py
@@ -0,0 +1,78 @@
+import logging
+import re
+import sys
+
+from sagemaker.deserializers import JSONDeserializer
+from sagemaker.pytorch import PyTorchModel
+from sagemaker.serverless import ServerlessInferenceConfig
+
+
+memory_size = { 
+    "dev": 4096 ,
+    "staging": 4096,
+    "prod": 6144 ,
+    "default": 4096,
+}
+max_concurrency = { 
+    "dev": 5,
+    "staging": 5,
+    "prod": 10,
+    "default": 5,
+}
+
+
+def deploy(
+    name: str,
+    stage: str,
+    version: str,
+    model_data: str,
+    role: str,
+):
+    sagemaker_logger = logging.getLogger("sagemaker")
+    sagemaker_logger.setLevel(logging.DEBUG)
+    sagemaker_logger.addHandler(logging.StreamHandler(sys.stdout))
+
+    composed_name =  re.sub(
+        r"[^a-zA-Z0-9\-]", "-", f"{name}-{version}-{stage}")
+
+    model = PyTorchModel(
+        name=composed_name,
+        model_data=model_data,
+        framework_version="1.12",
+        py_version="py38",
+        role=role,
+        env={
+            "SAGEMAKER_MODEL_SERVER_TIMEOUT": "3600",
+            "TS_MAX_RESPONSE_SIZE": "2000000000",
+            "TS_MAX_REQUEST_SIZE": "2000000000",
+            "MMS_MAX_RESPONSE_SIZE": "2000000000",
+            "MMS_MAX_REQUEST_SIZE": "2000000000",
+        },
+    )
+
+
+    return model.deploy(
+        initial_instance_count=1,
+        deserializer=JSONDeserializer(),
+        endpoint_name=composed_name,
+        serverless_inference_config=ServerlessInferenceConfig(
+            memory_size_in_mb=memory_size[stage],
+            max_concurrency=max_concurrency[stage]
+        )
+    )
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Deploy a model to Amazon SageMaker")
+
+    parser.add_argument("--name", type=str, required=True, help="Name of the model")
+    parser.add_argument("--stage", type=str, required=True, help="Stage of the model")
+    parser.add_argument("--version", type=str, required=True, help="Version of the model")
+    parser.add_argument("--model_data", type=str, required=True, help="S3 location of the model data")
+    parser.add_argument("--role", type=str, required=True, help="ARN of the IAM role to use")
+
+    args = parser.parse_args()
+
+    deploy(name=args.name, stage=args.stage, version=args.version, model_data=args.model_data, role=args.role)
diff --git a/example-get-started-experiments/code/src/endpoint_prediction.py b/example-get-started-experiments/code/src/endpoint_prediction.py
@@ -0,0 +1,56 @@
+from io import BytesIO
+from pathlib import Path
+
+import dvc.api
+import numpy as np
+from PIL import Image
+from sagemaker.deserializers import NumpyDeserializer
+from sagemaker.pytorch import PyTorchPredictor
+from sagemaker.serializers import IdentitySerializer
+
+
+def paint_mask(mask, color_map={0: (0, 0, 0), 1: (0, 0, 255)}):
+    vis_shape = mask.shape + (3,)
+    vis = np.zeros(vis_shape)
+    for i, c in color_map.items():
+        vis[mask == i] = color_map[i]
+    return Image.fromarray(vis.astype(np.uint8))
+
+
+def endpoint_prediction(
+    img_path: str,
+    endpoint_name: str,
+    output_path: str = "predictions",
+):
+    params = dvc.api.params_show()
+    img_size = params["train"]["img_size"]
+    predictor = PyTorchPredictor(endpoint_name, serializer=IdentitySerializer(), deserializer=NumpyDeserializer())
+    name = endpoint_name
+
+    output_file = Path(output_path) / name / Path(img_path).name
+    output_file.parent.mkdir(exist_ok=True, parents=True)
+
+    io = BytesIO()
+    Image.open(img_path).resize((img_size, img_size)).save(io, format="PNG")
+    result = predictor.predict(io.getvalue())[0]
+
+    img_pil = Image.open(img_path)
+    overlay_img_pil = Image.blend(
+        img_pil.convert("RGBA"), 
+        paint_mask(result).convert("RGBA").resize(img_pil.size), 
+        0.5
+    )
+    overlay_img_pil.save(str(output_file.with_suffix(".png")))
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description='Run inference on an image using a SageMaker endpoint')
+    parser.add_argument('--img_path', type=str, help='path to the input image')
+    parser.add_argument('--endpoint_name', type=str, help='name of the SageMaker endpoint to use')
+    parser.add_argument('--output_path', type=str, default='predictions', help='path to save the output predictions')
+
+    args = parser.parse_args()
+
+    endpoint_prediction(args.img_path, args.endpoint_name, args.output_path)
diff --git a/example-get-started-experiments/code/src/train.py b/example-get-started-experiments/code/src/train.py
@@ -68,6 +68,7 @@ def train():
         models_dir = Path("models")
         models_dir.mkdir(exist_ok=True)
         learn.export(fname=(models_dir / "model.pkl").absolute())
+        torch.save(learn.model, (models_dir / "model.pth").absolute())
         live.log_artifact(
             str(models_dir / "model.pkl"),
             type="model",

diff --git a/example-get-started-experiments/generate.sh b/example-get-started-experiments/generate.sh
@@ -90,11 +90,10 @@ git add .
 tick
 git commit -m "Run notebook and apply best experiment"
 git tag -a "1-notebook-dvclive" -m "Experiment using Notebook"
-gto register results/train:pool-segmentation --version v1.0.0
-gto assign results/train:pool-segmentation --version v1.0.0 --stage dev
 
 
 cp -r $HERE/code/src .
+cp -r $HERE/code/sagemaker .
 cp $HERE/code/params.yaml .
 sed -e "s/base_lr: 0.01/base_lr: $BEST_EXP_BASE_LR/" -i".bkp" params.yaml
 rm params.yaml.bkp
@@ -109,14 +108,18 @@ dvc remove models/model.pkl.dvc
 dvc stage add -n train \
   -p base,train \
   -d src/train.py -d data/train_data \
-  -o models/model.pkl \
+  -o models/model.pkl -o models/model.pth \
   python src/train.py
 
 dvc stage add -n evaluate \
   -p base,evaluate \
   -d src/evaluate.py -d models/model.pkl -d data/test_data \
   python src/evaluate.py
 
+dvc stage add -n sagemaker \
+  -d models/model.pth -o model.tar.gz \
+  'cp models/model.pth sagemaker/code/model.pth && cd sagemaker && tar -cpzf model.tar.gz code/ && cd .. && mv sagemaker/model.tar.gz .  && rm sagemaker/code/model.pth'
-  'cp models/model.pth sagemaker/code/model.pth && cd sagemaker && tar -cpzf model.tar.gz code/ && cd .. && mv sagemaker/model.tar.gz .  && rm sagemaker/code/model.pth'
+  'cp models/model.pth sagemaker/out && cp sagemaker/code sagemaker/out && tar -cpzf model.tar.gz sagemaker/out'
-  'cp models/model.pth sagemaker/code/model.pth && cd sagemaker && tar -cpzf model.tar.gz code/ && cd .. && mv sagemaker/model.tar.gz .  && rm sagemaker/code/model.pth'
+  'cp models/model.pth sagemaker/out && cp sagemaker/code sagemaker/out && tar -cpzf model.tar.gz sagemaker/out'
+
 git add .
 tick
 git commit -m "Convert Notebook to dvc.yaml pipeline"
@@ -127,9 +130,8 @@ git add .
 tick
 git commit -m "Run dvc.yaml pipeline"
 git tag -a "2-dvc-pipeline" -m "Experiment using dvc pipeline"
-gto register results/train:pool-segmentation --version v1.0.1
-gto assign results/train:pool-segmentation --version v1.0.0 --stage prod
-gto assign results/train:pool-segmentation --version v1.0.1 --stage dev
+gto register results/train:pool-segmentation --version v0.1.0
+gto assign results/train:pool-segmentation --version v0.1.0 --stage dev
 
 export GIT_AUTHOR_NAME="David de la Iglesia"
 export GIT_AUTHOR_EMAIL="daviddelaiglesiacastro@gmail.com"