refactor to a python module and add tests

.gitignore

@@ -35,3 +35,9 @@ notebooks/sampled2_dataframe_with_zeroshot_masks.csv
 notebooks/sampled_dataframe_with_zeroshot_masks.csv
 notebooks/segment.gpkg
 notebooks/tmp.py
+notebooks/burn_embeds.npz
+sam_decode_mar/
+scripts/save_npz_embeds.py
+src/sam_serve/__pycache__/
+tests/__pycache__/
+

Dockerfile-build

@@ -16,13 +16,17 @@ RUN pip install "onnxruntime==1.14.1" "onnx==1.13.1" "git+https://github.com/fac
 
 COPY handler*.py ./
 
-RUN torch-model-archiver --model-name sam_vit_h_encode --version 1.0.0 --serialized-file model-weights/sam_vit_h_4b8939.pth --handler handler_encode.py
+COPY requirements.txt ./requirements.txt
+
+RUN torch-model-archiver --model-name sam_vit_h_encode --version 1.0.0 \
+  --serialized-file model-weights/sam_vit_h_4b8939.pth --handler handler_encode.py
 
 COPY scripts ./scripts
 
 RUN python scripts/export_onnx_model.py --checkpoint model-weights/sam_vit_h_4b8939.pth --model-type vit_h --output sam_vit_h_decode.onnx
 
-RUN torch-model-archiver --model-name sam_vit_h_decode --version 1.0.0 --serialized-file sam_vit_h_decode.onnx --handler handler_decode.py
+RUN torch-model-archiver --model-name sam_vit_h_decode --version 1.0.0 \
+  --serialized-file sam_vit_h_decode.onnx --handler handler_decode.py
 
 RUN mkdir -p /home/model-store \
     && mv *.mar /home/model-store/ \

Dockerfile-cpu

@@ -1,4 +1,4 @@
-FROM pytorch/torchserve:latest-cpu
+FROM pytorch/torchserve:0.8.2-cpu
 
 USER root
 
@@ -8,7 +8,7 @@ RUN apt update \
     && unzip awscliv2.zip \
     && ./aws/install
 
-RUN pip install "onnxruntime==1.14.1" "onnx==1.13.1" "git+https://github.com/facebookresearch/segment-anything.git" "shapely" "pyproj" "rasterio"
+RUN pip install "onnxruntime==1.14.1" "onnx==1.13.1" "git+https://github.com/facebookresearch/segment-anything.git" "shapely" "pyproj" "rasterio" "debugpy"
 
 COPY ./deployment/config_decode.properties /home/model-server/config.properties
 COPY ./deployment/append_memory_setting.sh /home/model-server/append_memory_setting.sh
@@ -18,7 +18,9 @@ RUN chmod +x /home/model-server/append_memory_setting.sh \
 COPY start ./
 
 WORKDIR /home/model-server
+COPY dist/sam_serve-1.0.0-py3-none-any.whl /home/model-server
+RUN pip install /home/model-server/sam_serve-1.0.0-py3-none-any.whl
 ENV TEMP=/home/model-server/tmp
-ENV ENABLE_TORCH_PROFILER=TRUE
+ENV ENABLE_TORCH_PROFILER=FALSE
 
 CMD ./start

Dockerfile-gpu

@@ -1,4 +1,4 @@
-FROM pytorch/torchserve:latest-gpu
+FROM pytorch/torchserve:0.8.2-gpu
 
 USER root
 
@@ -18,7 +18,9 @@ RUN chmod +x /home/model-server/append_memory_setting.sh \
 COPY start ./
 
 WORKDIR /home/model-server
+COPY dist/sam_serve-1.0.0-py3-none-any.whl /home/model-server
+RUN pip install /home/model-server/sam_serve-1.0.0-py3-none-any.whl
 ENV TEMP=/home/model-server/volume/tmp
-ENV ENABLE_TORCH_PROFILER=TRUE
+ENV ENABLE_TORCH_PROFILER=FALSE
 
 CMD ./start

README.md

@@ -46,11 +46,47 @@ bash start_serve_decode_cpu.sh
 ```
 
 ### Test the encode service on the CPU
+
 The CPU service is served on 7080 by default. 8080 for the GPU service by default.
+
 ```
 curl http://127.0.0.1:7080/predictions/sam_vit_h_encode -T ./data/sample-img-fox.jpg
 ```
 
+### Testing
+
+All tests in `tests/` cover the functionality of the decoder. Logic in `decode.py` is run in pytest fixtures within `conftest.py` when outputs need to be shared by different tests in `test_decode.py`.
+
+To start running the tests, make sure you have the test models. You should have the same models used during inference, including
+
+```
+(test.py3.10) (base) rave@rave-desktop:~/segment-anything-services/tests/models$ tree
+.
+├── sam_vit_h_4b8939.pth
+└── sam_vit_h_decode.onnx
+```
+
+You can get both by unzipping the .mar archives copied to `model-store` from the sam-builder container you started in the previous step. Then, move the .onnx and the .pth files to ./tests/models/
+
+```
+unzip model-store/sam_vit_h_decode.mar -d ./sam_decode_mar
+unzip model-store/sam_vit_h_encode.mar -d ./sam_encode_mar
+cp ./sam_decode_mar/sam_vit_h_decode.onnx ./tests/models/
+cp ./sam_decode_mar/sam_vit_h_4b8939.pth ./tests/models/
+```
+
+Install the testing environment with [hatch](https://hatch.pypa.io/latest/install/): `pip install hatch`
+
+Then, create the environment. I tested with Python 3.10, Python 3.11 does not work because of an onnxruntime version issue.
+
+`hatch -e test.py3.10 shell`
+
+Then, run tests with pytest
+
+```
+pytest tests
+```
+
 ## Local Setup without Docker
 
 ### 1. Downloading model weights
@@ -63,6 +99,7 @@ aws s3 sync s3://segment-anything/model-weights/ model-weights
 
 otherwise, get checkpoints from the original repo: https://github.com/facebookresearch/segment-anything/tree/main#model-checkpoints
 
+
 ### 2a. Package the torch weights for GPU encoding
 
 This step takes a long time presumably because the uncompiled weights are massive. Packaging the ONNX model is faster in the later steps.
@@ -115,6 +152,27 @@ docker run -it --rm \
     --gpus all sam-dev
 ```
 
+## Deployment
+
+1. Install the dependencies to deploy from npm
+
+```
+npm install -g @openaddresses/deploy
+```
+2. Set up your credentials by navigating to the AWS SSO login page, and selecting "Command Line Access" to copy the temporary credentials. paste these to `~/.aws/credentials` and rename the AWS profile (ask Ryan for this.)
+3. `touch ~/.deployrc.json` and fill it with
+```
+   {
+    "profilename": {
+        "region": "us-east-1"
+    }
+}
+```
+4. run deploy init to generate a `.deployrc.json` config for the repo. both the local and global config are needed to deploy. Fill in args when prompted based on the profile name
+5. commit and make a PR with any changes. wait for all github actions to complete so that the model archives and docker images are built
+6. `deploy update prod` to deploy changes to prod
+
+
 ### (Potentially) Frequently Asked Questions
 Q: Why two services?
 

deployment/config_decode.properties

@@ -10,3 +10,4 @@ vmargs=-Xmx8g
 cors_allowed_origin=*
 cors_allowed_methods=GET, POST, PUT, OPTIONS
 cors_allowed_headers=X-Custom-Header,content-type,content-length,crossorigin
+install_py_dep_per_model=true

deployment/config_encode.properties

@@ -11,3 +11,4 @@ max_response_size=17797905
 cors_allowed_origin=*
 cors_allowed_methods=GET, POST, PUT, OPTIONS
 cors_allowed_headers=X-Custom-Header,content-type,content-length,crossorigin
+install_py_dep_per_model=true