Add an example of using Hamilton in AWS SageMaker

Author: tyapochkin

examples/aws/sagemaker/README.md

@@ -0,0 +1,74 @@
+# Deploying Hamilton Functions as an AWS SageMaker Processing Job
+
+[AWS SageMaker](https://aws.amazon.com/sagemaker/) is a comprehensive platform that facilitates the creation, training, and deployment of machine learning (ML) models. This guide demonstrates deploying a "hello-world" [processing job](https://docs.aws.amazon.com/sagemaker/latest/dg/processing-job.html) using Hamilton functions on SageMaker.
+
+## Prerequisites
+
+- **AWS CLI Setup**: Ensure that the AWS CLI is configured on your machine. Follow the [Quick Start guide](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-quickstart.html) for setup instructions.
+
+## Step-by-Step Guide
+
+### 1. Build the Docker Image
+
+Navigate to the container directory and build the Docker image:
+
+```shell
+cd container/ && docker build --platform linux/amd64 -t aws-sagemaker-hamilton . && cd ..
+```
+
+### 2. Create AWS ECR repository.
+
+Replace `111122223333` with your AWS account number.
+
+- **Authenticate Docker to Amazon ECR**:
+
+    Retrieve an authentication token to authenticate your Docker client to your Amazon Elastic Container Registry (ECR):
+
+    ```shell
+    aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 111122223333.dkr.ecr.us-east-1.amazonaws.com
+    ```
+
+- **Create the ECR Repository**:
+
+    ```shell
+    aws ecr create-repository --repository-name aws-sagemaker-hamilton --region us-east-1 --image-scanning-configuration scanOnPush=true --image-tag-mutability MUTABLE
+    ```
+
+### 3. Deploy the Image to AWS ECR
+
+Ensure the AWS account number is correctly replaced with yours:
+
+```shell
+docker tag aws-sagemaker-hamilton 111122223333.dkr.ecr.us-east-1.amazonaws.com/aws-sagemaker-hamilton:latest
+docker push 111122223333.dkr.ecr.us-east-1.amazonaws.com/aws-sagemaker-hamilton:latest
+```
+
+### 4. Create simple role for AWS SageMaker ScriptProcessor.
+
+- **Create the Role**:
+
+    Example of creating an AWS Role with full permissions for ECR and S3.
+
+    ```shell
+    aws iam create-role --role-name SageMakerScriptProcessorRole --assume-role-policy-document '{"Version": "2012-10-17", "Statement": [{ "Effect": "Allow", "Principal": { "Service": "sagemaker.amazonaws.com"}, "Action": "sts:AssumeRole"}]}'
+    ```
+
+- **Attach Policies to the Role**:
+
+    Here we grant full access to ECR, S3 and SageMaker as an example. For production environments it's important to restrict access appropriately.
+
+    ```shell
+    aws iam attach-role-policy --role-name SageMakerScriptProcessorRole --policy-arn arn:aws:iam::aws:policy/AmazonS3FullAccess
+    aws iam attach-role-policy --role-name SageMakerScriptProcessorRole --policy-arn arn:aws:iam::aws:policy/AmazonEC2ContainerRegistryFullAccess
+    aws iam attach-role-policy --role-name SageMakerScriptProcessorRole --policy-arn arn:aws:iam::aws:policy/AmazonSageMakerFullAccess
+    ```
+
+### 5. Install additional requirements
+
+```shell
+pip install -r requirements.txt
+```
+
+### 6. Execute the Processing Job
+
+Find the detailed example in [notebook.ipynb](notebook.ipynb) to run the processing job.

examples/aws/sagemaker/app/__init__.py

@@ -0,0 +1,38 @@
+import pandas as pd
+
+from hamilton.function_modifiers import extract_columns
+
+
+@extract_columns("spend", "signups")
+def raw_table(input_table: pd.DataFrame) -> pd.DataFrame:
+    return input_table
+
+
+def avg_3wk_spend(spend: pd.Series) -> pd.Series:
+    """Rolling 3 week average spend."""
+    return spend.rolling(3).mean()
+
+
+def spend_per_signup(spend: pd.Series, signups: pd.Series) -> pd.Series:
+    """The cost per signup in relation to spend."""
+    return spend / signups
+
+
+def spend_mean(spend: pd.Series) -> float:
+    """Shows function creating a scalar. In this case it computes the mean of the entire column."""
+    return spend.mean()
+
+
+def spend_zero_mean(spend: pd.Series, spend_mean: float) -> pd.Series:
+    """Shows function that takes a scalar. In this case to zero mean spend."""
+    return spend - spend_mean
+
+
+def spend_std_dev(spend: pd.Series) -> float:
+    """Function that computes the standard deviation of the spend column."""
+    return spend.std()
+
+
+def spend_zero_mean_unit_variance(spend_zero_mean: pd.Series, spend_std_dev: float) -> pd.Series:
+    """Function showing one way to make spend have zero mean and unit variance."""
+    return spend_zero_mean / spend_std_dev

examples/aws/sagemaker/app/functions.py

@@ -0,0 +1,11 @@
+FROM python:3.11-slim-buster
+
+RUN apt-get update && apt-get install -y graphviz
+
+COPY requirements.txt ./
+
+RUN pip install -r requirements.txt
+
+ENV HAMILTON_TELEMETRY_ENABLED=false
+
+ENTRYPOINT ["python3"]

examples/aws/sagemaker/container/Dockerfile

@@ -0,0 +1,2 @@
+pandas
+sf-hamilton[visualization]

examples/aws/sagemaker/container/requirements.txt

@@ -0,0 +1,7 @@
+signups,spend
+1,10
+10,10
+50,20
+100,40
+200,40
+400,50

examples/aws/sagemaker/data/input_table.csv

@@ -0,0 +1,86 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from sagemaker.processing import ScriptProcessor, ProcessingInput, ProcessingOutput"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "script_processor = ScriptProcessor(\n",
+    "    command=['python3'],\n",
+    "    image_uri='<account_number>.dkr.ecr.<region>.amazonaws.com/aws-sagemaker-hamilton:latest',  # Change to the actual URI\n",
+    "    role='arn:aws:iam::<account_number>:role/SageMakerScriptProcessorRole',  # Change to the actual URI\n",
+    "    instance_count=1,\n",
+    "    instance_type='ml.t3.medium'\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# All inputs and outputs should be located in /opt/ml/processing/\n",
+    "script_processor.run(\n",
+    "    code='processing.py',\n",
+    "    inputs=[\n",
+    "        ProcessingInput(\n",
+    "            source='data/',\n",
+    "            destination='/opt/ml/processing/input/data'\n",
+    "        ),\n",
+    "        ProcessingInput(\n",
+    "            source='app/',\n",
+    "            destination='/opt/ml/processing/input/code/app'\n",
+    "        )\n",
+    "    ],\n",
+    "    outputs=[\n",
+    "        ProcessingOutput(\n",
+    "            source='/opt/ml/processing/output/',\n",
+    "            destination='s3://path/to/output/directory'  # Change to the actual URI\n",
+    "        )\n",
+    "    ]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "After job finishes, the new files will appear in `s3://path/to/output/directory`:\n",
+    "- `output_table.csv`\n",
+    "- `dag_visualization.svg`"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "hamilton",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

examples/aws/sagemaker/notebook.ipynb

@@ -0,0 +1,29 @@
+import pandas as pd
+from app import functions
+
+from hamilton import driver
+
+if __name__ == "__main__":
+
+    df = pd.read_csv("/opt/ml/processing/input/data/input_table.csv")
+
+    dr = driver.Driver({}, functions)
+
+    inputs = {"input_table": df}
+
+    output_columns = [
+        "spend",
+        "signups",
+        "avg_3wk_spend",
+        "spend_per_signup",
+        "spend_zero_mean_unit_variance",
+    ]
+
+    # DAG visualization
+    dot = dr.visualize_execution(final_vars=output_columns, inputs=inputs)
+    with open("/opt/ml/processing/output/dag_visualization.svg", "wb") as svg_out:
+        svg_out.write(dot.pipe(format="svg"))
+
+    # DAG execution
+    df_result = dr.execute(output_columns, inputs=inputs)
+    df_result.to_csv("/opt/ml/processing/output/output_table.csv")

examples/aws/sagemaker/processing.py

@@ -0,0 +1 @@
+sagemaker