├── .gitignore
├── Dockerfile
├── LICENSE
├── Makefile
├── README.md
├── app
├── auth
│ └── callback
│ │ └── route.ts
├── globals.css
├── layout.tsx
├── login
│ ├── page.tsx
│ └── submit-button.tsx
├── page.tsx
└── protected
│ └── page.tsx
├── components
├── AuthButton.tsx
└── Header.tsx
├── configs
└── huggingface-tc-bert-base-cased-202403291810.yaml
├── example_configs
├── bert-eqa.yaml
├── bert-tc.yaml
├── llama7b.yaml
└── train-bert.yaml
├── middleware.ts
├── model_manager.py
├── next.config.js
├── package-lock.json
├── package.json
├── pdm.lock
├── postcss.config.js
├── pyproject.toml
├── requirements.txt
├── scripts
└── setup_role.sh
├── server.py
├── setup.sh
├── src
├── __init__.py
├── config.py
├── console.py
├── huggingface
│ ├── __init__.py
│ └── hf_hub_api.py
├── main.py
├── sagemaker
│ ├── __init__.py
│ ├── create_model.py
│ ├── delete_model.py
│ ├── fine_tune_model.py
│ ├── query_endpoint.py
│ ├── resources.py
│ └── search_jumpstart_models.py
├── schemas
│ ├── __init__.py
│ ├── deployment.py
│ ├── model.py
│ ├── query.py
│ └── training.py
├── session.py
├── utils
│ ├── __init__.py
│ ├── aws_utils.py
│ ├── format.py
│ ├── model_utils.py
│ └── rich_utils.py
└── yaml.py
├── tailwind.config.js
├── tests
├── __init__.py
└── test_openai_proxy.py
├── tsconfig.json
└── utils
└── supabase
├── client.ts
├── middleware.ts
└── server.ts
/.gitignore:
--------------------------------------------------------------------------------
1 | # See https://help.github.com/articles/ignoring-files/ for more about ignoring files.
2 |
3 | # dependencies
4 | /node_modules
5 | /.pnp
6 | .pnp.js
7 | .yarn/install-state.gz
8 |
9 | # testing
10 | /coverage
11 |
12 | # next.js
13 | /.next/
14 | /out/
15 |
16 | # production
17 | /build
18 |
19 | # misc
20 | .DS_Store
21 | *.pem
22 |
23 | # debug
24 | npm-debug.log*
25 | yarn-debug.log*
26 | yarn-error.log*
27 |
28 | # local env files
29 | .env*.local
30 |
31 | # vercel
32 | .vercel
33 |
34 | # typescript
35 | *.tsbuildinfo
36 | next-env.d.ts
37 |
38 | # Byte-compiled / optimized / DLL files
39 | __pycache__/
40 | *.py[cod]
41 | *$py.class
42 |
43 | # C extensions
44 | *.so
45 |
46 | # Distribution / packaging
47 | .Python
48 | build/
49 | develop-eggs/
50 | dist/
51 | downloads/
52 | eggs/
53 | .eggs/
54 | lib/
55 | lib64/
56 | parts/
57 | sdist/
58 | var/
59 | wheels/
60 | share/python-wheels/
61 | *.egg-info/
62 | .installed.cfg
63 | *.egg
64 | MANIFEST
65 |
66 | # PyInstaller
67 | # Usually these files are written by a python script from a template
68 | # before PyInstaller builds the exe, so as to inject date/other infos into it.
69 | *.manifest
70 | *.spec
71 |
72 | # Installer logs
73 | pip-log.txt
74 | pip-delete-this-directory.txt
75 |
76 | # Unit test / coverage reports
77 | htmlcov/
78 | .tox/
79 | .nox/
80 | .coverage
81 | .coverage.*
82 | .cache
83 | nosetests.xml
84 | coverage.xml
85 | *.cover
86 | *.py,cover
87 | .hypothesis/
88 | .pytest_cache/
89 | cover/
90 |
91 | # Translations
92 | *.mo
93 | *.pot
94 |
95 | # Django stuff:
96 | *.log
97 | local_settings.py
98 | db.sqlite3
99 | db.sqlite3-journal
100 |
101 | # Flask stuff:
102 | instance/
103 | .webassets-cache
104 |
105 | # Scrapy stuff:
106 | .scrapy
107 |
108 | # Sphinx documentation
109 | docs/_build/
110 |
111 | # PyBuilder
112 | .pybuilder/
113 | target/
114 |
115 | # Jupyter Notebook
116 | .ipynb_checkpoints
117 |
118 | # IPython
119 | profile_default/
120 | ipython_config.py
121 |
122 | # pyenv
123 | # For a library or package, you might want to ignore these files since the code is
124 | # intended to run in multiple environments; otherwise, check them in:
125 | # .python-version
126 |
127 | # pipenv
128 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
129 | # However, in case of collaboration, if having platform-specific dependencies or dependencies
130 | # having no cross-platform support, pipenv may install dependencies that don't work, or not
131 | # install all needed dependencies.
132 | #Pipfile.lock
133 |
134 | # poetry
135 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
136 | # This is especially recommended for binary packages to ensure reproducibility, and is more
137 | # commonly ignored for libraries.
138 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
139 | #poetry.lock
140 |
141 | # pdm
142 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
143 | #pdm.lock
144 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
145 | # in version control.
146 | # https://pdm.fming.dev/#use-with-ide
147 | .pdm.toml
148 |
149 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
150 | __pypackages__/
151 |
152 | # Celery stuff
153 | celerybeat-schedule
154 | celerybeat.pid
155 |
156 | # SageMath parsed files
157 | *.sage.py
158 |
159 | # Environments
160 | .env
161 | .venv
162 | env/
163 | venv/
164 | ENV/
165 | env.bak/
166 | venv.bak/
167 |
168 | # Spyder project settings
169 | .spyderproject
170 | .spyproject
171 |
172 | # Rope project settings
173 | .ropeproject
174 |
175 | # mkdocs documentation
176 | /site
177 |
178 | # mypy
179 | .mypy_cache/
180 | .dmypy.json
181 | dmypy.json
182 |
183 | # Pyre type checker
184 | .pyre/
185 |
186 | # pytype static type analyzer
187 | .pytype/
188 |
189 | # Cython debug symbols
190 | cython_debug/
191 |
192 | # PyCharm
193 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
194 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
195 | # and can be added to the global gitignore or merged into this file. For a more nuclear
196 | # option (not recommended) you can uncomment the following to ignore the entire idea folder.
197 | #.idea/
198 |
199 | # vscode
200 | .vscode/
201 | # es-lint
202 | .eslintrc.cjs
203 |
204 | models/
205 | configs/
206 | .pdm-python
207 |
--------------------------------------------------------------------------------
/Dockerfile:
--------------------------------------------------------------------------------
1 | # Use an official Python runtime as a parent image
2 | FROM python:3.12-slim
3 |
4 | # Set environment variables
5 | ENV PYTHONDONTWRITEBYTECODE 1
6 | ENV PYTHONUNBUFFERED 1
7 |
8 | # Set the working directory in the container
9 | WORKDIR /app
10 |
11 | # Copy the project files to the working directory
12 | COPY . .
13 |
14 | # Install PDM and use it to install dependencies
15 | RUN pip install --no-cache-dir pdm \
16 | && pdm install --no-interactive
17 |
18 | # Expose port 8000 to the outside world
19 | EXPOSE 8000
20 |
21 | # Run uvicorn when the container launches
22 | CMD ["pdm", "run", "uvicorn", "server:app", "--host", "0.0.0.0", "--port", "8000"]
23 |
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 | MIT License
2 |
3 | Copyright (c) 2024 OpenFoundry
4 |
5 | Permission is hereby granted, free of charge, to any person obtaining a copy
6 | of this software and associated documentation files (the "Software"), to deal
7 | in the Software without restriction, including without limitation the rights
8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 |
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 |
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 |
--------------------------------------------------------------------------------
/Makefile:
--------------------------------------------------------------------------------
1 | venv: venv/touchfile
2 |
3 | venv/touchfile: requirements.txt
4 | test -d venv || python3 -m venv venv
5 | . venv/bin/activate; pip install -Ur requirements.txt
6 | touch venv/touchfile
7 | clean:
8 | rm -rf venv
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 |
2 |
7 |
Model Manager v0.1, by
8 |
9 | 
10 |
11 |
12 | Deploy open source AI models to AWS in minutes.
13 |
15 |
21 | Table of Contents
22 |
23 |
24 | About Model Manager
25 |
26 |
27 | Getting Started
28 |
32 |
33 | Using Model Manager What we're working on next Known issues Contributing License Contact
40 |
41 |
42 |
43 | ## About Model Manager
44 | Model Manager is a Python tool that simplifies the process of deploying an open source AI model to your own cloud. Instead of spending hours digging through documentation to figure out how to get AWS working, Model Manager lets you deploy open source AI models directly from the command line.
45 |
46 | Choose a model from Hugging Face or SageMaker, and Model Manager will spin up a SageMaker instance with a ready-to-query endpoint in minutes.
47 |
48 | Here we’re deploying Microsoft’s Phi2. Larger models such as this take about 10 minutes to spin up.
49 |
50 | https://github.com/openfoundry-ai/model_manager/assets/164248540/f7fbf9ce-04c3-49b0-b5be-2977b5cc90af
51 |
52 | (back to top )
59 |
60 |
61 |
62 | (back to top )
110 |
111 |
112 |
113 |
114 | (back to top )
211 |
212 |
213 |
214 |
215 | (back to top )
224 |
225 |
226 |
227 | (back to top )
238 |
239 |
240 |
241 | (back to top )
250 |
251 |
252 |
253 |
254 | (back to top )
261 |
262 |
263 |
264 |
265 |
23 | {children}
24 |
25 |
26 |
27 | );
28 | }
29 |
--------------------------------------------------------------------------------
/app/login/page.tsx:
--------------------------------------------------------------------------------
1 | import Link from "next/link";
2 | import { headers } from "next/headers";
3 | import { createClient } from "@/utils/supabase/server";
4 | import { redirect } from "next/navigation";
5 | import { SubmitButton } from "./submit-button";
6 |
7 | export default function Login({
8 | searchParams,
9 | }: {
10 | searchParams: { message: string };
11 | }) {
12 | const signIn = async (formData: FormData) => {
13 | "use server";
14 |
15 | const email = formData.get("email") as string;
16 | const password = formData.get("password") as string;
17 | const supabase = createClient();
18 |
19 | const { error } = await supabase.auth.signInWithPassword({
20 | email,
21 | password,
22 | });
23 |
24 | if (error) {
25 | return redirect("/login?message=Could not authenticate user");
26 | }
27 |
28 | return redirect("/protected");
29 | };
30 |
31 | const signUp = async (formData: FormData) => {
32 | "use server";
33 |
34 | const origin = headers().get("origin");
35 | const email = formData.get("email") as string;
36 | const password = formData.get("password") as string;
37 | const supabase = createClient();
38 |
39 | const { error } = await supabase.auth.signUp({
40 | email,
41 | password,
42 | options: {
43 | emailRedirectTo: `${origin}/auth/callback`,
44 | },
45 | });
46 |
47 | if (error) {
48 | return redirect("/login?message=Could not authenticate user");
49 | }
50 |
51 | return redirect("/login?message=Check email to continue sign in process");
52 | };
53 |
54 | return (
55 |
56 |
60 |
72 | {" "}
74 | Back
75 |
76 |
77 |
117 |
17 | {isPending ? pendingText : children}
18 |
19 | );
20 | }
21 |
--------------------------------------------------------------------------------
/app/page.tsx:
--------------------------------------------------------------------------------
1 | import AuthButton from "../components/AuthButton";
2 | import { createClient } from "@/utils/supabase/server";
3 |
4 | export default async function Index() {
5 | const canInitSupabaseClient = () => {
6 | // This function is just for the interactive tutorial.
7 | // Feel free to remove it once you have Supabase connected.
8 | try {
9 | createClient();
10 | return true;
11 | } catch (e) {
12 | return false;
13 | }
14 | };
15 |
16 | const isSupabaseConnected = canInitSupabaseClient();
17 |
18 | return (
19 |
20 |
21 |
22 | Model Manager
23 | {isSupabaseConnected &&
}
24 |
26 |
27 |
40 |
19 |
20 |
21 | This is a protected page that you can only see as an authenticated
22 | user
23 |
24 |
25 |
26 | Model Manager
27 |
28 |
30 |
31 |
32 |
33 |
34 |
35 |
36 |
49 |
22 | Hey, {user.email}!
23 |
28 |
29 | ) : (
30 | /tmp/assume-role-policy-document.json
17 | {
18 | "Version": "2012-10-17",
19 | "Statement": [{
20 | "Effect": "Allow",
21 | "Principal": {
22 | "Service": "sagemaker.amazonaws.com"
23 | },
24 | "Action": "sts:AssumeRole"
25 | }]
26 | }
27 | EOF
28 |
29 | # Creates the role
30 | echo "SAGEMAKER_ROLE=`aws iam create-role --role-name ${ROLE_NAME} --assume-role-policy-document file:///tmp/assume-role-policy-document.json | grep -Eo '"arn:aws:iam.*?[^\\]"'`" >> .env
31 |
32 | # attaches the Sagemaker full access policy to the role
33 | aws iam attach-role-policy --policy-arn ${POLICY} --role-name ${ROLE_NAME}
--------------------------------------------------------------------------------
/server.py:
--------------------------------------------------------------------------------
1 | import uvicorn
2 | import os
3 | from dotenv import dotenv_values
4 | from fastapi import FastAPI
5 | from src.config import get_config_for_endpoint, get_endpoints_for_model
6 | from src.sagemaker.resources import get_sagemaker_endpoint
7 | from src.sagemaker.query_endpoint import make_query_request
8 | from src.schemas.query import Query, ChatCompletion
9 | from src.session import session
10 | from litellm import completion
11 |
12 | os.environ["AWS_REGION_NAME"] = session.region_name
13 | app = FastAPI()
14 |
15 |
16 | class NotDeployedException(Exception):
17 | pass
18 |
19 |
20 | @app.get("/endpoint/{endpoint_name}")
21 | def get_endpoint(endpoint_name: str):
22 | return get_sagemaker_endpoint(endpoint_name)
23 |
24 |
25 | @app.post("/endpoint/{endpoint_name}/query")
26 | def query_endpoint(endpoint_name: str, query: Query):
27 | config = get_config_for_endpoint(endpoint_name)
28 | if query.context is None:
29 | query.context = ''
30 |
31 | # Support multi-model endpoints
32 | config = (config.deployment, config.models[0])
33 | return make_query_request(endpoint_name, query, config)
34 |
35 |
36 | @app.post("/chat/completions")
37 | def chat_completion(chat_completion: ChatCompletion):
38 | model_id = chat_completion.model
39 |
40 | # Validate model is for completion tasks
41 | endpoints = get_endpoints_for_model(model_id)
42 | if len(endpoints) == 0:
43 | raise NotDeployedException
44 |
45 | messages = chat_completion.messages
46 | # Currently using the first available endpoint.
47 | endpoint_name = endpoints[0].deployment.endpoint_name
48 |
49 | res = completion(
50 | model=f"sagemaker/{endpoint_name}",
51 | messages=messages,
52 | temperature=0.9,
53 | hf_model_name=model_id,
54 | )
55 |
56 | return res
57 |
58 |
59 | if __name__ == "__main__":
60 | uvicorn.run(app, host="0.0.0.0", port=8000)
61 |
--------------------------------------------------------------------------------
/setup.sh:
--------------------------------------------------------------------------------
1 | #!/bin/sh
2 | python3 -m ensurepip --upgrade
3 | make
4 | source venv/bin/activate
5 | if ! command -v aws &> /dev/null
6 | then
7 | OS="$(uname -s)"
8 | case "${OS}" in
9 | Linux*)
10 | curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"
11 | unzip awscliv2.zip
12 | sudo ./aws/install
13 | ;;
14 | Darwin*)
15 | curl "https://awscli.amazonaws.com/AWSCLIV2.pkg" -o "AWSCLIV2.pkg"
16 | sudo installer -pkg AWSCLIV2.pkg -target /
17 | ;;
18 | *)
19 | echo "Unsupported OS: ${OS}. See https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html"
20 | exit 1
21 | ;;
22 | esac
23 | fi
24 | aws configure set region us-east-1 && aws configure
25 | touch .env
26 | if ! grep -q "SAGEMAKER_ROLE" .env
27 | then
28 | bash ./scripts/setup_role.sh
29 | fi
30 | source venv/bin/activate
--------------------------------------------------------------------------------
/src/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/openfoundry-ai/model_manager/34f9ffc3e1962aa87ca088203535d8d0506b4e7b/src/__init__.py
--------------------------------------------------------------------------------
/src/config.py:
--------------------------------------------------------------------------------
1 | import glob
2 | import yaml
3 | from src.yaml import dumper
4 | from src.schemas.model import Model
5 | from src.schemas.deployment import Deployment
6 | from typing import Dict, Tuple, List, NamedTuple, Optional
7 |
8 |
9 | class ModelDeployment(NamedTuple):
10 | deployment: Deployment
11 | models: List[Model]
12 |
13 |
14 | def get_deployment_configs(path: Optional[str] = None) -> List[ModelDeployment]:
15 | if path is None:
16 | path = "./configs/*.yaml"
17 |
18 | configurations = glob.glob(path)
19 | configs = []
20 |
21 | for configuration in configurations:
22 | with open(configuration) as config:
23 | configuration = yaml.safe_load(config)
24 | if configuration is None:
25 | continue
26 |
27 | # Filter out training configs
28 | if configuration.get('deployment') is None:
29 | continue
30 |
31 | deployment = configuration['deployment']
32 | models = configuration['models']
33 | configs.append(ModelDeployment(
34 | deployment=deployment, models=models))
35 |
36 | return configs
37 |
38 |
39 | def get_endpoints_for_model(model_id: str, path: Optional[str] = None) -> List[ModelDeployment]:
40 | configs = get_deployment_configs(path)
41 | endpoints = []
42 | for config in configs:
43 | models = [model.id for model in config.models]
44 | if model_id in models:
45 | # TODO: Check if endpoint is still active
46 | endpoints.append(config)
47 | return endpoints
48 |
49 |
50 | def get_config_for_endpoint(endpoint_name: str) -> Optional[ModelDeployment]:
51 | configs = get_deployment_configs()
52 | for config in configs:
53 | if config.deployment.endpoint_name == endpoint_name:
54 | return config
55 | return None
56 |
57 |
58 | def write_config(deployment: Deployment, model: Model):
59 | with open(f"./configs/{deployment.endpoint_name}.yaml", 'w') as config:
60 | out = {
61 | "deployment": deployment,
62 | "models": [model],
63 | }
64 | config.write(yaml.dump(out, Dumper=dumper))
65 | return
66 |
--------------------------------------------------------------------------------
/src/console.py:
--------------------------------------------------------------------------------
1 | from rich.console import Console
2 | console = Console()
3 |
--------------------------------------------------------------------------------
/src/huggingface/__init__.py:
--------------------------------------------------------------------------------
1 | from enum import StrEnum
2 | from huggingface_hub import HfApi
3 |
4 |
5 | class HuggingFaceTask(StrEnum):
6 | AudioClassification = "audio-classification"
7 | AutomaticSpeechRecognition = "automatic-speech-recognition"
8 | Conversational = "conversational"
9 | DepthEstimation = "depth-estimation"
10 | DocumentQuestionAnswering = "document-question-answering"
11 | FeatureExtraction = "feature-extraction"
12 | FillMask = "fill-mask"
13 | ImageClassification = "image-classification"
14 | ImageFeatureExtraction = "image-feature-extraction"
15 | ImageSegmentation = "image-segmentation"
16 | ImageToImage = "image-to-image"
17 | ImageToText = "image-to-text"
18 | MaskGeneration = "mask-generation"
19 | ObjectDetection = "object-detection"
20 | QuestionAnswering = "question-answering"
21 | Summarization = "summarization"
22 | TableQuestionAnswering = "table-question-answering"
23 | Text2TextGeneration = "text2text-generation"
24 | TextClassification = "text-classification"
25 | TextGeneration = "text-generation"
26 | TextToAudio = "text-to-audio"
27 | TokenClassification = "token-classification"
28 | Translation = "translation"
29 | TranslationXXtoYY = "translation_xx_to_yy"
30 | VideoClassification = "video-classification"
31 | VisualQuestionAnswering = "visual-question-answering"
32 | ZeroShotClassification = "zero-shot-classification"
33 | ZeroShotImageClassification = "zero-shot-image-classification"
34 | ZeroShotAudioClassification = "zero-shot-audio-classification"
35 | ZeroShotObjectDetection = "zero-shot-object-detection"
36 |
37 |
38 | AVAILABLE_PIPELINES = ["feature-extraction",
39 | "text-classification", "token-classification", "question-answering",
40 | "table-question-answering", "fill-mask", "summarization", "translation",
41 | "text2text-generation", "text-generation", "zero-shot-classification", "conversational",
42 | "image-classification", "translation_XX_to_YY"]
43 |
44 | hf_api = HfApi()
45 |
--------------------------------------------------------------------------------
/src/huggingface/hf_hub_api.py:
--------------------------------------------------------------------------------
1 | from dotenv import dotenv_values
2 | from src.console import console
3 | from src.huggingface import hf_api
4 | from src.schemas.model import Model
5 | from src.utils.rich_utils import print_error
6 |
7 | HUGGING_FACE_HUB_TOKEN = dotenv_values(".env").get("HUGGING_FACE_HUB_KEY")
8 |
9 |
10 | def get_hf_task(model: Model):
11 | task = None
12 | try:
13 | model_info = hf_api.model_info(
14 | model.id, token=HUGGING_FACE_HUB_TOKEN)
15 | task = model_info.pipeline_tag
16 | if model_info.transformers_info is not None and model_info.transformers_info.pipeline_tag is not None:
17 | task = model_info.transformers_info.pipeline_tag
18 | except Exception:
19 | # better error handling for auth
20 | console.print_exception()
21 | print_error("Model not found, please try another.")
22 | return
23 | return task
24 |
--------------------------------------------------------------------------------
/src/main.py:
--------------------------------------------------------------------------------
1 | import inquirer
2 | import logging
3 | import threading
4 | from InquirerPy import prompt
5 | from src.sagemaker import EC2Instance
6 | from src.sagemaker.create_model import deploy_model
7 | from src.sagemaker.delete_model import delete_sagemaker_model
8 | from src.sagemaker.resources import list_sagemaker_endpoints, select_instance, list_service_quotas_async
9 | from src.sagemaker.query_endpoint import make_query_request
10 | from src.sagemaker.search_jumpstart_models import search_sagemaker_jumpstart_model
11 | from src.utils.rich_utils import print_error, print_success
12 | from src.schemas.deployment import Deployment, Destination
13 | from src.schemas.model import Model, ModelSource
14 | from src.schemas.query import Query
15 | from src.config import get_config_for_endpoint
16 | from enum import StrEnum
17 | from rich import print
18 |
19 |
20 | class Actions(StrEnum):
21 | LIST = "Show active model endpoints"
22 | DEPLOY = "Deploy a model endpoint"
23 | DELETE = "Delete a model endpoint"
24 | QUERY = "Query a model endpoint"
25 | EXIT = "Quit"
26 | # TRAIN = "fine tune a model"
27 |
28 |
29 | def main(args, loglevel):
30 | logging.basicConfig(format="%(levelname)s: %(message)s", level=loglevel)
31 |
32 | print("[magenta]Model Manager by OpenFoundry.")
33 | print("[magenta]Star us on Github ☆! [blue]https://github.com/openfoundry-ai/model_manager")
34 |
35 | # list_service_quotas is a pretty slow API and it's paginated.
36 | # Use async here and store the result in instances
37 | instances = []
38 | instance_thread = threading.Thread(
39 | target=list_service_quotas_async, args=[instances])
40 | instance_thread.start()
41 |
42 | while True:
43 | active_endpoints = list_sagemaker_endpoints()
44 | questions = [
45 | inquirer.List(
46 | 'action',
47 | message="What would you like to do?",
48 | choices=[e.value for e in Actions]
49 | )
50 | ]
51 | answers = inquirer.prompt(questions)
52 | if (answers is None):
53 | break
54 |
55 | action = answers['action']
56 |
57 | match action:
58 | case Actions.LIST:
59 | if len(active_endpoints) != 0:
60 | [print(f"[blue]{endpoint['EndpointName']}[/blue] running on an [green]{endpoint['InstanceType']}[/green] instance")
61 | for endpoint in active_endpoints]
62 | print('\n')
63 | else:
64 | print_error('No active endpoints.\n')
65 | case Actions.DEPLOY:
66 | build_and_deploy_model(instances, instance_thread)
67 | case Actions.DELETE:
68 | if (len(active_endpoints) == 0):
69 | print_success("No Endpoints to delete!")
70 | continue
71 |
72 | questions = [
73 | inquirer.Checkbox('endpoints',
74 | message="Which endpoints would you like to delete? (space to select)",
75 | choices=[endpoint['EndpointName']
76 | for endpoint in active_endpoints]
77 | )
78 | ]
79 | answers = inquirer.prompt(questions)
80 | if (answers is None):
81 | continue
82 |
83 | endpoints_to_delete = answers['endpoints']
84 | delete_sagemaker_model(endpoints_to_delete)
85 | case Actions.QUERY:
86 | if (len(active_endpoints) == 0):
87 | print_success("No Endpoints to query!")
88 | continue
89 |
90 | questions = [
91 | {
92 | "type": "list",
93 | "message": "Which endpoint would you like to query?",
94 | "name": 'endpoint',
95 | "choices": [endpoint['EndpointName'] for endpoint in active_endpoints]
96 | },
97 | {
98 | "type": "input",
99 | "message": "What would you like to query?",
100 | "name": 'query'
101 | }
102 | ]
103 | answers = prompt(questions)
104 | if (answers is None):
105 | continue
106 |
107 | endpoint = answers['endpoint']
108 | query = Query(query=answers['query'])
109 | config = get_config_for_endpoint(endpoint)
110 |
111 | # support multi-model endpoints
112 | config = (config.deployment, config.models[0])
113 | make_query_request(endpoint, query, config)
114 | case Actions.EXIT:
115 | quit()
116 |
117 | print_success("Goodbye!")
118 |
119 |
120 | class ModelType(StrEnum):
121 | SAGEMAKER = "Deploy a Sagemaker model"
122 | HUGGINGFACE = "Deploy a Hugging Face model"
123 | CUSTOM = "Deploy a custom model"
124 |
125 |
126 | def build_and_deploy_model(instances, instance_thread):
127 | questions = [
128 | inquirer.List(
129 | 'model_type',
130 | message="Choose a model type:",
131 | choices=[model_type.value for model_type in ModelType]
132 | )
133 | ]
134 | answers = inquirer.prompt(questions)
135 | if answers is None:
136 | return
137 |
138 | model_type = answers['model_type']
139 |
140 | match model_type:
141 | case ModelType.SAGEMAKER:
142 | models = []
143 | while len(models) == 0:
144 | models = search_sagemaker_jumpstart_model()
145 | if models is None:
146 | quit()
147 |
148 | questions = [
149 | {
150 | "type": "fuzzy",
151 | "name": "model_id",
152 | "message": "Choose a model. Search by task (e.g. eqa) or model name (e.g. llama):",
153 | "choices": models,
154 | "match_exact": True,
155 | }
156 | ]
157 | answers = prompt(questions)
158 | if (answers is None):
159 | return
160 |
161 | model_id, model_version = answers['model_id'], "2.*"
162 | instance_thread.join()
163 | instance_type = select_instance(instances)
164 |
165 | model = Model(
166 | id=model_id,
167 | model_version=model_version,
168 | source=ModelSource.Sagemaker
169 | )
170 |
171 | deployment = Deployment(
172 | instance_type=instance_type,
173 | destination=Destination.AWS,
174 | num_gpus=4 if instance_type == EC2Instance.LARGE else 1
175 | )
176 |
177 | predictor = deploy_model(deployment=deployment, model=model)
178 | case ModelType.HUGGINGFACE:
179 | questions = [
180 | inquirer.Text(
181 | 'model_id',
182 | message="Enter the exact model name from huggingface.co (e.g. google-bert/bert-base-uncased)",
183 | )
184 | ]
185 | answers = inquirer.prompt(questions)
186 | if answers is None:
187 | return
188 |
189 | model_id = answers['model_id']
190 | instance_thread.join()
191 | instance_type = select_instance(instances)
192 |
193 | model = Model(
194 | id=model_id,
195 | source=ModelSource.HuggingFace
196 | )
197 |
198 | deployment = Deployment(
199 | instance_type=instance_type,
200 | destination=Destination.AWS,
201 | )
202 |
203 | predictor = deploy_model(deployment=deployment, model=model)
204 |
205 | case ModelType.CUSTOM:
206 | questions = [
207 | {
208 | "type": "input",
209 | "message": "What is the local path or S3 URI of the model?",
210 | "name": "path"
211 | },
212 | {
213 | "type": "input",
214 | "message": "What is the base model that was fine tuned? (e.g. google-bert/bert-base-uncased)",
215 | "name": "base_model"
216 | }
217 | ]
218 | answers = prompt(questions)
219 | local_path = answers['path']
220 | base_model = answers['base_model']
221 |
222 | instance_thread.join()
223 | instance_type = select_instance(instances)
224 |
225 | model = Model(
226 | id=base_model,
227 | source=ModelSource.Custom,
228 | location=local_path
229 | )
230 |
231 | deployment = Deployment(
232 | instance_type=instance_type,
233 | destination=Destination.AWS
234 | )
235 | deploy_model(deployment=deployment, model=model)
236 |
--------------------------------------------------------------------------------
/src/sagemaker/__init__.py:
--------------------------------------------------------------------------------
1 | from enum import StrEnum
2 |
3 |
4 | class EC2Instance(StrEnum):
5 | SMALL = "ml.m5.xlarge"
6 | MEDIUM = "ml.p3.2xlarge"
7 | LARGE = "ml.g5.12xlarge"
8 |
9 |
10 | class SagemakerTask(StrEnum):
11 | AutomaticSpeechRecognition = "asr"
12 | AudioEmbedding = "audioembedding"
13 | Classification = "classification"
14 | Depth2img = "depth2img"
15 | ExtractiveQuestionAnswering = "eqa"
16 | FillMask = "fillmask"
17 | ImageClassification = "ic"
18 | ImageEmbedding = "icembedding"
19 | ImageGeneration = "imagegeneration"
20 | Inpainting = "inpainting"
21 | ImageSegmentation = "is"
22 | LLM = "llm"
23 | NamedEntityRecognition = "ner"
24 | ObjectDetection = "od"
25 | Od1 = "od1"
26 | Regression = "regression"
27 | SemanticSegmentation = "semseg"
28 | SentenceSimilarity = "sentencesimilarity"
29 | SentencePairClassification = "spc"
30 | Summarization = "summarization"
31 | TabTransformerClassification = "tabtransformerclassification"
32 | TabTransformerRegression = "tabtransformerregression"
33 | TextClassification = "tc"
34 | TcEmbedding = "tcembedding"
35 | Text2text = "text2text"
36 | TextEmbedding = "textembedding"
37 | TextGeneration = "textgeneration"
38 | TextGeneration1 = "textgeneration1"
39 | TextGeneration2 = "textgeneration2"
40 | TextGenerationJP = "textgenerationjp"
41 | TextGenerationNeuron = "textgenerationneuron"
42 | Translation = "translation"
43 | Txt2img = "txt2img"
44 | Upscaling = "upscaling"
45 | ZeroShotTextClassification = "zstc"
46 |
47 | @staticmethod
48 | def list():
49 | return list(map(lambda c: c.value, SagemakerTask))
50 |
--------------------------------------------------------------------------------
/src/sagemaker/create_model.py:
--------------------------------------------------------------------------------
1 | import json
2 | from dotenv import dotenv_values
3 | from rich.table import Table
4 | from sagemaker import image_uris, model_uris, script_uris
5 | from sagemaker.huggingface import get_huggingface_llm_image_uri
6 | from sagemaker.huggingface.model import HuggingFaceModel
7 | from sagemaker.jumpstart.model import JumpStartModel
8 | from sagemaker.jumpstart.estimator import JumpStartEstimator
9 | from sagemaker.model import Model
10 | from sagemaker.predictor import Predictor
11 | from sagemaker.s3 import S3Uploader
12 | from src.config import write_config
13 | from src.schemas.model import Model, ModelSource
14 | from src.schemas.deployment import Deployment
15 | from src.session import session, sagemaker_session
16 | from src.console import console
17 | from src.utils.aws_utils import construct_s3_uri, is_s3_uri
18 | from src.utils.rich_utils import print_error, print_success
19 | from src.utils.model_utils import get_unique_endpoint_name, get_model_and_task
20 | from src.huggingface import HuggingFaceTask
21 | from src.huggingface.hf_hub_api import get_hf_task
22 |
23 | HUGGING_FACE_HUB_TOKEN = dotenv_values(".env").get("HUGGING_FACE_HUB_KEY")
24 | SAGEMAKER_ROLE = dotenv_values(".env")["SAGEMAKER_ROLE"]
25 |
26 |
27 | # TODO: Consolidate
28 | def deploy_model(deployment: Deployment, model: Model):
29 | match model.source:
30 | case ModelSource.HuggingFace:
31 | deploy_huggingface_model(deployment, model)
32 | case ModelSource.Sagemaker:
33 | create_and_deploy_jumpstart_model(deployment, model)
34 | case ModelSource.Custom:
35 | deploy_custom_huggingface_model(deployment, model)
36 |
37 |
38 | def deploy_huggingface_model(deployment: Deployment, model: Model):
39 | region_name = session.region_name
40 | task = get_hf_task(model)
41 | model.task = task
42 | env = {
43 | 'HF_MODEL_ID': model.id,
44 | 'HF_TASK': task,
45 | }
46 |
47 | if HUGGING_FACE_HUB_TOKEN is not None:
48 | env['HUGGING_FACE_HUB_TOKEN'] = HUGGING_FACE_HUB_TOKEN
49 |
50 | image_uri = None
51 | if deployment.num_gpus:
52 | env['SM_NUM_GPUS'] = json.dumps(deployment.num_gpus)
53 |
54 | if deployment.quantization:
55 | env['HF_MODEL_QUANTIZE'] = deployment.quantization
56 |
57 | if task == HuggingFaceTask.TextGeneration:
58 | # use TGI imageq if llm.
59 | image_uri = get_huggingface_llm_image_uri(
60 | "huggingface",
61 | version="1.4.2"
62 | )
63 |
64 | huggingface_model = HuggingFaceModel(
65 | env=env,
66 | role=SAGEMAKER_ROLE,
67 | transformers_version="4.37",
68 | pytorch_version="2.1",
69 | py_version="py310",
70 | image_uri=image_uri
71 | )
72 |
73 | endpoint_name = get_unique_endpoint_name(
74 | model.id, deployment.endpoint_name)
75 |
76 | deployment.endpoint_name = endpoint_name
77 |
78 | console.log(
79 | "Deploying model to AWS. [magenta]This may take up to 10 minutes for very large models.[/magenta] See full logs here:")
80 | console.print(
81 | f"https://{region_name}.console.aws.amazon.com/cloudwatch/home#logsV2:log-groups/log-group/$252Faws$252Fsagemaker$252FEndpoints$252F{endpoint_name}")
82 |
83 | with console.status("[bold green]Deploying model...") as status:
84 | table = Table(show_header=False, header_style="magenta")
85 | table.add_column("Resource", style="dim")
86 | table.add_column("Value", style="blue")
87 | table.add_row("model", model.id)
88 | table.add_row("EC2 instance type", deployment.instance_type)
89 | table.add_row("Number of instances", str(
90 | deployment.instance_count))
91 | table.add_row("task", task)
92 | console.print(table)
93 |
94 | try:
95 | predictor = huggingface_model.deploy(
96 | initial_instance_count=deployment.instance_count,
97 | instance_type=deployment.instance_type,
98 | endpoint_name=endpoint_name,
99 | )
100 | except Exception:
101 | console.print_exception()
102 | quit()
103 |
104 | print_success(
105 | f"{model.id} is now up and running at the endpoint [blue]{predictor.endpoint_name}")
106 |
107 | write_config(deployment, model)
108 | return predictor
109 |
110 |
111 | def deploy_custom_huggingface_model(deployment: Deployment, model: Model):
112 | region_name = session.region_name
113 | if model.location is None:
114 | print_error("Missing model source location.")
115 | return
116 |
117 | s3_path = model.location
118 | if not is_s3_uri(model.location):
119 | # Local file. Upload to s3 before deploying
120 | bucket = sagemaker_session.default_bucket()
121 | s3_path = construct_s3_uri(bucket, f"models/{model.id}")
122 | with console.status(f"[bold green]Uploading custom {model.id} model to S3 at {s3_path}...") as status:
123 | try:
124 | s3_path = S3Uploader.upload(
125 | model.location, s3_path)
126 | except Exception:
127 | print_error("[red] Model failed to upload to S3")
128 |
129 | endpoint_name = get_unique_endpoint_name(
130 | model.id, deployment.endpoint_name)
131 |
132 | deployment.endpoint_name = endpoint_name
133 | model.task = get_model_and_task(model.id)['task']
134 |
135 | console.log(
136 | "Deploying model to AWS. [magenta]This may take up to 10 minutes for very large models.[/magenta] See full logs here:")
137 | console.print(
138 | f"https://{region_name}.console.aws.amazon.com/cloudwatch/home#logsV2:log-groups/log-group/$252Faws$252Fsagemaker$252FEndpoints$252F{endpoint_name}")
139 |
140 | # create Hugging Face Model Class
141 | huggingface_model = HuggingFaceModel(
142 | # path to your trained sagemaker model
143 | model_data=s3_path,
144 | role=SAGEMAKER_ROLE, # iam role with permissions to create an Endpoint
145 | transformers_version="4.37",
146 | pytorch_version="2.1",
147 | py_version="py310",
148 | )
149 |
150 | with console.status("[bold green]Deploying model...") as status:
151 | table = Table(show_header=False, header_style="magenta")
152 | table.add_column("Resource", style="dim")
153 | table.add_column("Value", style="blue")
154 | table.add_row("S3 Path", s3_path)
155 | table.add_row("EC2 instance type", deployment.instance_type)
156 | table.add_row("Number of instances", str(
157 | deployment.instance_count))
158 | console.print(table)
159 |
160 | try:
161 | predictor = huggingface_model.deploy(
162 | initial_instance_count=deployment.instance_count,
163 | instance_type=deployment.instance_type,
164 | endpoint_name=endpoint_name
165 | )
166 | except Exception:
167 | console.print_exception()
168 | quit()
169 |
170 | print_success(
171 | f"Custom {model.id} is now up and running at the endpoint [blue]{predictor.endpoint_name}")
172 |
173 | write_config(deployment, model)
174 | return predictor
175 |
176 |
177 | def create_and_deploy_jumpstart_model(deployment: Deployment, model: Model):
178 | region_name = session.region_name
179 | endpoint_name = get_unique_endpoint_name(
180 | model.id, deployment.endpoint_name)
181 | deployment.endpoint_name = endpoint_name
182 | model.task = get_model_and_task(model.id)['task']
183 |
184 | console.log(
185 | "Deploying model to AWS. [magenta]This may take up to 10 minutes for very large models.[/magenta] See full logs here:")
186 |
187 | console.print(
188 | f"https://{region_name}.console.aws.amazon.com/cloudwatch/home#logsV2:log-groups/log-group/$252Faws$252Fsagemaker$252FEndpoints$252F{endpoint_name}")
189 |
190 | with console.status("[bold green]Deploying model...") as status:
191 | table = Table(show_header=False, header_style="magenta")
192 | table.add_column("Resource", style="dim")
193 | table.add_column("Value", style="blue")
194 | table.add_row("model", model.id)
195 | table.add_row("EC2 instance type", deployment.instance_type)
196 | table.add_row("Number of instances", str(
197 | deployment.instance_count))
198 | console.print(table)
199 |
200 | jumpstart_model = JumpStartModel(
201 | model_id=model.id, instance_type=deployment.instance_type, role=SAGEMAKER_ROLE)
202 |
203 | # Attempt to deploy to AWS
204 | try:
205 | predictor = jumpstart_model.deploy(
206 | initial_instance_count=deployment.instance_count,
207 | instance_type=deployment.instance_type,
208 | endpoint_name=endpoint_name,
209 | accept_eula=True
210 | )
211 | pass
212 | except Exception:
213 | console.print_exception()
214 | quit()
215 |
216 | write_config(deployment, model)
217 | print_success(
218 | f"{model.id} is now up and running at the endpoint [blue]{predictor.endpoint_name}")
219 |
220 | return predictor
221 |
--------------------------------------------------------------------------------
/src/sagemaker/delete_model.py:
--------------------------------------------------------------------------------
1 | import boto3
2 | from rich import print
3 | from src.utils.rich_utils import print_success
4 | from typing import List
5 |
6 |
7 | def delete_sagemaker_model(endpoint_names: List[str] = None):
8 | sagemaker_client = boto3.client('sagemaker')
9 |
10 | if len(endpoint_names) == 0:
11 | print_success("No Endpoints to delete!")
12 | return
13 |
14 | # Add validation / error handling
15 | for endpoint in endpoint_names:
16 | print(f"Deleting [blue]{endpoint}")
17 | sagemaker_client.delete_endpoint(EndpointName=endpoint)
18 |
--------------------------------------------------------------------------------
/src/sagemaker/fine_tune_model.py:
--------------------------------------------------------------------------------
1 | import logging
2 | import os
3 | import sagemaker
4 | from botocore.exceptions import ClientError
5 | from datasets import load_dataset
6 | from rich import print
7 | from rich.table import Table
8 | from sagemaker.jumpstart.estimator import JumpStartEstimator
9 | from src.console import console
10 | from src.schemas.model import Model, ModelSource
11 | from src.schemas.training import Training
12 | from src.session import sagemaker_session
13 | from src.utils.aws_utils import is_s3_uri
14 | from src.utils.rich_utils import print_success, print_error
15 | from transformers import AutoTokenizer
16 |
17 | from dotenv import load_dotenv
18 | load_dotenv()
19 | SAGEMAKER_ROLE = os.environ.get("SAGEMAKER_ROLE")
20 |
21 |
22 | def prep_hf_data(s3_bucket: str, dataset_name_or_path: str, model: Model):
23 | train_dataset, test_dataset = load_dataset(
24 | dataset_name_or_path, split=["train", "test"])
25 | tokenizer = AutoTokenizer.from_pretrained(model.id)
26 |
27 | def tokenize(batch):
28 | return tokenizer(batch["text"], padding="max_length", truncation=True)
29 |
30 | # tokenize train and test datasets
31 | train_dataset = train_dataset.map(tokenize, batched=True)
32 | test_dataset = test_dataset.map(tokenize, batched=True)
33 |
34 | # set dataset format for PyTorch
35 | train_dataset = train_dataset.rename_column("label", "labels")
36 | train_dataset.set_format(
37 | "torch", columns=["input_ids", "attention_mask", "labels"])
38 | test_dataset = test_dataset.rename_column("label", "labels")
39 | test_dataset.set_format(
40 | "torch", columns=["input_ids", "attention_mask", "labels"])
41 |
42 | # save train_dataset to s3
43 | training_input_path = f's3://{s3_bucket}/datasets/train'
44 | train_dataset.save_to_disk(training_input_path)
45 |
46 | # save test_dataset to s3
47 | test_input_path = f's3://{s3_bucket}/datasets/test'
48 | test_dataset.save_to_disk(test_input_path)
49 |
50 | return training_input_path, test_input_path
51 |
52 |
53 | def train_model(training: Training, model: Model, estimator):
54 | # TODO: Accept hf datasets or local paths to upload to s3
55 | if not is_s3_uri(training.training_input_path):
56 | raise Exception("Training data needs to be uploaded to s3")
57 |
58 | # TODO: Implement training, validation, and test split or accept a directory of files
59 | training_dataset_s3_path = training.training_input_path
60 |
61 | table = Table(show_header=False, header_style="magenta")
62 | table.add_column("Resource", style="dim")
63 | table.add_column("Value", style="blue")
64 | table.add_row("model", model.id)
65 | table.add_row("model_version", model.version)
66 | table.add_row("base_model_uri", estimator.model_uri)
67 | table.add_row("image_uri", estimator.image_uri)
68 | table.add_row("EC2 instance type", training.instance_type)
69 | table.add_row("Number of instances", str(training.instance_count))
70 | console.print(table)
71 |
72 | estimator.fit({"training": training_dataset_s3_path})
73 |
74 | predictor = estimator.deploy(
75 | initial_instance_count=training.instance_count, instance_type=training.instance_type)
76 |
77 | print_success(
78 | f"Trained model {model.id} is now up and running at the endpoint [blue]{predictor.endpoint_name}")
79 |
80 |
81 | def fine_tune_model(training: Training, model: Model):
82 | estimator = None
83 | match model.source:
84 | case ModelSource.Sagemaker:
85 | hyperparameters = get_hyperparameters_for_model(training, model)
86 | estimator = JumpStartEstimator(
87 | model_id=model.id,
88 | model_version=model.version,
89 | instance_type=training.instance_type,
90 | instance_count=training.instance_count,
91 | output_path=training.output_path,
92 | environment={"accept_eula": "true"},
93 | role=SAGEMAKER_ROLE,
94 | sagemaker_session=sagemaker_session,
95 | hyperparameters=hyperparameters
96 | )
97 | case ModelSource.HuggingFace:
98 | raise NotImplementedError
99 | case ModelSource.Custom:
100 | raise NotImplementedError
101 |
102 | try:
103 | print_success("Enqueuing training job")
104 | res = train_model(training, model, estimator)
105 | except ClientError as e:
106 | logging.error(e)
107 | print_error("Training job enqueue fail")
108 | return False
109 |
110 |
111 | def get_hyperparameters_for_model(training: Training, model: Model):
112 | hyperparameters = sagemaker.hyperparameters.retrieve_default(
113 | model_id=model.id, model_version=model.version)
114 |
115 | if training.hyperparameters is not None:
116 | hyperparameters.update(
117 | (k, v) for k, v in training.hyperparameters.model_dump().items() if v is not None)
118 | return hyperparameters
119 |
--------------------------------------------------------------------------------
/src/sagemaker/query_endpoint.py:
--------------------------------------------------------------------------------
1 | import boto3
2 | import json
3 | import inquirer
4 | from InquirerPy import prompt
5 | from sagemaker.huggingface.model import HuggingFacePredictor
6 | from src.config import ModelDeployment
7 | from src.console import console
8 | from src.sagemaker import SagemakerTask
9 | from src.huggingface import HuggingFaceTask
10 | from src.utils.model_utils import get_model_and_task, is_sagemaker_model, get_text_generation_hyperpameters
11 | from src.utils.rich_utils import print_error
12 | from src.schemas.deployment import Deployment
13 | from src.schemas.model import Model
14 | from src.schemas.query import Query
15 | from src.session import sagemaker_session
16 | from typing import Dict, Tuple, Optional
17 |
18 |
19 | def make_query_request(endpoint_name: str, query: Query, config: Tuple[Deployment, Model]):
20 | if is_sagemaker_model(endpoint_name, config):
21 | return query_sagemaker_endpoint(endpoint_name, query, config)
22 | else:
23 | return query_hugging_face_endpoint(endpoint_name, query, config)
24 |
25 |
26 | def parse_response(query_response):
27 | model_predictions = json.loads(query_response['Body'].read())
28 | probabilities, labels, predicted_label = model_predictions[
29 | 'probabilities'], model_predictions['labels'], model_predictions['predicted_label']
30 | return probabilities, labels, predicted_label
31 |
32 |
33 | def query_hugging_face_endpoint(endpoint_name: str, user_query: Query, config: Tuple[Deployment, Model]):
34 | task = get_model_and_task(endpoint_name, config)['task']
35 | predictor = HuggingFacePredictor(endpoint_name=endpoint_name,
36 | sagemaker_session=sagemaker_session)
37 |
38 | query = user_query.query
39 | context = user_query.context
40 |
41 | input = {"inputs": query}
42 | if task is not None and task == HuggingFaceTask.QuestionAnswering:
43 | if context is None:
44 | questions = [{
45 | "type": "input", "message": "What context would you like to provide?:", "name": "context"}]
46 | answers = prompt(questions)
47 | context = answers.get('context', '')
48 |
49 | if not context:
50 | raise Exception("Must provide context for question-answering")
51 |
52 | input = {}
53 | input['context'] = answers['context']
54 | input['question'] = query
55 |
56 | if task is not None and task == HuggingFaceTask.TextGeneration:
57 | parameters = get_text_generation_hyperpameters(config, user_query)
58 | input['parameters'] = parameters
59 |
60 | if task is not None and task == HuggingFaceTask.ZeroShotClassification:
61 | if context is None:
62 | questions = [
63 | inquirer.Text('labels',
64 | message="What labels would you like to use? (comma separated values)?",
65 | )
66 | ]
67 | answers = inquirer.prompt(questions)
68 | context = answers.get('labels', '')
69 |
70 | if not context:
71 | raise Exception(
72 | "Must provide labels for zero shot text classification")
73 |
74 | labels = context.split(',')
75 | input = json.dumps({
76 | "sequences": query,
77 | "candidate_labels": labels
78 | })
79 |
80 | try:
81 | result = predictor.predict(input)
82 | except Exception:
83 | console.print_exception()
84 | quit()
85 |
86 | print(result)
87 | return result
88 |
89 |
90 | def query_sagemaker_endpoint(endpoint_name: str, user_query: Query, config: Tuple[Deployment, Model]):
91 | client = boto3.client('runtime.sagemaker')
92 | task = get_model_and_task(endpoint_name, config)['task']
93 |
94 | if task not in [
95 | SagemakerTask.ExtractiveQuestionAnswering,
96 | SagemakerTask.TextClassification,
97 | SagemakerTask.SentenceSimilarity,
98 | SagemakerTask.SentencePairClassification,
99 | SagemakerTask.Summarization,
100 | SagemakerTask.NamedEntityRecognition,
101 | SagemakerTask.TextEmbedding,
102 | SagemakerTask.TcEmbedding,
103 | SagemakerTask.TextGeneration,
104 | SagemakerTask.TextGeneration1,
105 | SagemakerTask.TextGeneration2,
106 | SagemakerTask.Translation,
107 | SagemakerTask.FillMask,
108 | SagemakerTask.ZeroShotTextClassification
109 | ]:
110 | print_error("""
111 | Querying this model type inside of Model Manager isn’t yet supported.
112 | You can query it directly through the API endpoint - see here for documentation on how to do this:
113 | https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_runtime_InvokeEndpoint.html
114 | """)
115 | raise Exception("Unsupported")
116 |
117 | # MIME content type varies per deployment
118 | content_type = "application/x-text"
119 | accept_type = "application/json;verbose"
120 |
121 | # Depending on the task, input needs to be formatted differently.
122 | # e.g. question-answering needs to have {question: , context: }
123 | query = user_query.query
124 | context = user_query.context
125 | input = query.encode("utf-8")
126 | match task:
127 | case SagemakerTask.ExtractiveQuestionAnswering:
128 | if context is None:
129 | questions = [
130 | {
131 | 'type': 'input',
132 | 'name': 'context',
133 | 'message': "What context would you like to provide?",
134 | }
135 | ]
136 | answers = prompt(questions)
137 | context = answers.get("context", '')
138 |
139 | if not context:
140 | raise Exception("Must provide context for question-answering")
141 |
142 | content_type = "application/list-text"
143 | input = json.dumps([query, context]).encode("utf-8")
144 |
145 | case SagemakerTask.SentencePairClassification:
146 | if context is None:
147 | questions = [
148 | inquirer.Text('context',
149 | message="What sentence would you like to compare against?",
150 | )
151 | ]
152 | answers = inquirer.prompt(questions)
153 | context = answers.get("context", '')
154 | if not context:
155 | raise Exception(
156 | "Must provide a second sentence for sentence pair classification")
157 |
158 | content_type = "application/list-text"
159 | input = json.dumps([query, context]).encode("utf-8")
160 | case SagemakerTask.ZeroShotTextClassification:
161 | if context is None:
162 | questions = [
163 | inquirer.Text('labels',
164 | message="What labels would you like to use? (comma separated values)?",
165 | )
166 | ]
167 | answers = inquirer.prompt(questions)
168 | context = answers.get('labels', '')
169 |
170 | if not context:
171 | raise Exception(
172 | "must provide labels for zero shot text classification")
173 | labels = context.split(',')
174 |
175 | content_type = "application/json"
176 | input = json.dumps({
177 | "sequences": query,
178 | "candidate_labels": labels,
179 | }).encode("utf-8")
180 | case SagemakerTask.TextGeneration:
181 | parameters = get_text_generation_hyperpameters(config, user_query)
182 | input = json.dumps({
183 | "inputs": query,
184 | "parameters": parameters,
185 | }).encode("utf-8")
186 | content_type = "application/json"
187 |
188 | try:
189 | response = client.invoke_endpoint(
190 | EndpointName=endpoint_name, ContentType=content_type, Body=input, Accept=accept_type)
191 | except Exception:
192 | console.print_exception()
193 | quit()
194 |
195 | model_predictions = json.loads(response['Body'].read())
196 | print(model_predictions)
197 | return model_predictions
198 |
199 |
200 | def test(endpoint_name: str):
201 | text1 = 'astonishing ... ( frames ) profound ethical and philosophical questions in the form of dazzling pop entertainment'
202 | text2 = 'simply stupid , irrelevant and deeply , truly , bottomlessly cynical '
203 |
204 | for text in [text1, text2]:
205 | query_sagemaker_endpoint(endpoint_name, text.encode('utf-8'))
206 |
--------------------------------------------------------------------------------
/src/sagemaker/resources.py:
--------------------------------------------------------------------------------
1 | import boto3
2 | from functools import lru_cache
3 | from InquirerPy import inquirer
4 | from src.console import console
5 | from src.sagemaker import EC2Instance
6 | from src.config import get_config_for_endpoint
7 | from src.utils.format import format_sagemaker_endpoint, format_python_dict
8 | from typing import List, Tuple, Dict, Optional
9 |
10 |
11 | def list_sagemaker_endpoints(filter_str: str = None) -> List[str]:
12 | sagemaker_client = boto3.client('sagemaker')
13 |
14 | endpoints = sagemaker_client.list_endpoints()['Endpoints']
15 | if filter_str is not None:
16 | endpoints = list(filter(lambda x: filter_str ==
17 | x['EndpointName'], endpoints))
18 |
19 | for endpoint in endpoints:
20 | endpoint_config = sagemaker_client.describe_endpoint_config(
21 | EndpointConfigName=endpoint['EndpointName'])['ProductionVariants'][0]
22 | endpoint['InstanceType'] = endpoint_config['InstanceType']
23 | return endpoints
24 |
25 |
26 | def get_sagemaker_endpoint(endpoint_name: str) -> Optional[Dict[str, Optional[Dict]]]:
27 | endpoints = list_sagemaker_endpoints(endpoint_name)
28 | if not endpoints:
29 | return None
30 |
31 | endpoint = format_sagemaker_endpoint(endpoints[0])
32 |
33 | config = get_config_for_endpoint(endpoint_name)
34 | if config is None:
35 | return {'deployment': endpoint, 'model': None}
36 |
37 | deployment = format_python_dict(config.deployment.model_dump())
38 | formatted_models = []
39 | for model in config.models:
40 | model = format_python_dict(model.model_dump())
41 | formatted_models.append(model)
42 |
43 | # Merge the endpoint dict with our config
44 | deployment = {**endpoint, **deployment}
45 |
46 | return {
47 | 'deployment': deployment,
48 | 'models': formatted_models,
49 | }
50 |
51 |
52 | @lru_cache
53 | def list_service_quotas() -> List[Tuple[str, int]]:
54 | """ Gets a list of EC2 instances for inference """
55 |
56 | client = boto3.client('service-quotas')
57 | quotas = []
58 | try:
59 | response = client.list_service_quotas(
60 | ServiceCode="sagemaker",
61 | MaxResults=100,
62 | )
63 | next_token = response.get('NextToken')
64 | quotas = response['Quotas']
65 | while next_token is not None:
66 | response = client.list_service_quotas(
67 | ServiceCode="sagemaker",
68 | NextToken=next_token,
69 | )
70 | quotas.extend(response['Quotas'])
71 | next_token = response.get('NextToken')
72 | except client.exceptions.AccessDeniedException as error:
73 | console.print(
74 | "[red]User does not have access to Service Quotas. Grant access via IAM to get the list of available instances")
75 | return []
76 |
77 | # TODO: Filter for different usages
78 | available_instances = list(filter(lambda x: 'endpoint usage' in x[0] and x[1] > 0, [
79 | (quota['QuotaName'], quota['Value']) for quota in quotas]))
80 |
81 | # Clean up quota names
82 | available_instances = [(instance[0].split(" ")[0], instance[1])
83 | for instance in available_instances]
84 | return available_instances
85 |
86 |
87 | def list_service_quotas_async(instances=[]):
88 | """ Wrapper to allow access to list in threading """
89 | instances = instances.extend(list_service_quotas())
90 | return instances
91 |
92 |
93 | def select_instance(available_instances=None):
94 | choices = [instance[0] for instance in available_instances] or [
95 | instance for instance in EC2Instance]
96 | instance = inquirer.fuzzy(
97 | message="Choose an instance size (note: ml.m5.xlarge available by default; you must request quota from AWS to use other instance types):",
98 | choices=choices,
99 | default="ml.m5."
100 | ).execute()
101 | if instance is None:
102 | return
103 | return instance
104 |
--------------------------------------------------------------------------------
/src/sagemaker/search_jumpstart_models.py:
--------------------------------------------------------------------------------
1 | import inquirer
2 | from enum import StrEnum, auto
3 | from sagemaker.jumpstart.notebook_utils import list_jumpstart_models
4 | from src.utils.rich_utils import print_error
5 | from src.session import session, sagemaker_session
6 |
7 |
8 | class Frameworks(StrEnum):
9 | huggingface = auto()
10 | meta = auto()
11 | model = auto()
12 | tensorflow = auto()
13 | pytorch = auto()
14 | # autogluon
15 | # catboost
16 | # lightgbm
17 | mxnet = auto()
18 | # sklearn
19 | # xgboost
20 |
21 |
22 | def search_sagemaker_jumpstart_model():
23 | questions = [
24 | inquirer.List('framework',
25 | message="Which framework would you like to use?",
26 | choices=[framework.value for framework in Frameworks]
27 | ),
28 | ]
29 | answers = inquirer.prompt(questions)
30 | if answers is None:
31 | return
32 | filter_value = "framework == {}".format(answers["framework"])
33 |
34 | models = list_jumpstart_models(filter=filter_value,
35 | region=session.region_name, sagemaker_session=sagemaker_session)
36 | return models
37 |
--------------------------------------------------------------------------------
/src/schemas/__init__.py:
--------------------------------------------------------------------------------
1 | import src.schemas.deployment
2 | import src.schemas.model
3 | import src.schemas.query
4 | import src.schemas.training
5 |
--------------------------------------------------------------------------------
/src/schemas/deployment.py:
--------------------------------------------------------------------------------
1 | import yaml
2 | from src.yaml import loader, dumper
3 | from typing import Optional
4 | from enum import StrEnum
5 | from pydantic import BaseModel
6 |
7 |
8 | class Destination(StrEnum):
9 | AWS = "aws"
10 | # AZURE = "azure"
11 | # GCP = "gcp"
12 |
13 |
14 | class Deployment(BaseModel):
15 | destination: Destination
16 | instance_type: str
17 | endpoint_name: Optional[str] = None
18 | instance_count: Optional[int] = 1
19 | num_gpus: Optional[int] = None
20 | quantization: Optional[str] = None
21 |
22 |
23 | def deployment_representer(dumper: yaml.SafeDumper, deployment: Deployment) -> yaml.nodes.MappingNode:
24 | return dumper.represent_mapping("!Deployment", {
25 | "destination": deployment.destination.value,
26 | "instance_type": deployment.instance_type,
27 | "endpoint_name": deployment.endpoint_name,
28 | "instance_count": deployment.instance_count,
29 | "num_gpus": deployment.num_gpus,
30 | "quantization": deployment.quantization,
31 | })
32 |
33 |
34 | def deployment_constructor(loader: yaml.SafeLoader, node: yaml.nodes.ScalarNode) -> Deployment:
35 | return Deployment(**loader.construct_mapping(node))
36 |
37 |
38 | dumper.add_representer(Deployment, deployment_representer)
39 | loader.add_constructor(u'!Deployment', deployment_constructor)
40 |
--------------------------------------------------------------------------------
/src/schemas/model.py:
--------------------------------------------------------------------------------
1 | import yaml
2 | from src.yaml import loader, dumper
3 | from typing import Optional, Union, Dict
4 | from enum import StrEnum
5 | from src.huggingface import HuggingFaceTask
6 | from src.sagemaker import SagemakerTask
7 | from pydantic import BaseModel
8 |
9 |
10 | class ModelSource(StrEnum):
11 | HuggingFace = "huggingface"
12 | Sagemaker = "sagemaker"
13 | Custom = "custom"
14 |
15 |
16 | class Model(BaseModel):
17 | id: str
18 | source: ModelSource
19 | task: Optional[Union[HuggingFaceTask, SagemakerTask]] = None
20 | version: Optional[str] = None
21 | location: Optional[str] = None
22 | predict: Optional[Dict[str, str]] = None
23 |
24 |
25 | def model_representer(dumper: yaml.SafeDumper, model: Model) -> yaml.nodes.MappingNode:
26 | return dumper.represent_mapping("!Model", {
27 | "id": model.id,
28 | "source": model.source.value,
29 | "task": model.task,
30 | "version": model.version,
31 | "location": model.location,
32 | "predict": model.predict,
33 | })
34 |
35 |
36 | def model_constructor(loader: yaml.SafeLoader, node: yaml.nodes.MappingNode) -> Model:
37 | return Model(**loader.construct_mapping(node))
38 |
39 |
40 | dumper.add_representer(Model, model_representer)
41 | loader.add_constructor(u'!Model', model_constructor)
42 |
--------------------------------------------------------------------------------
/src/schemas/query.py:
--------------------------------------------------------------------------------
1 |
2 | from pydantic import BaseModel
3 | from typing import Optional, List, Literal
4 |
5 |
6 | class QueryParameters(BaseModel):
7 | max_length: Optional[int] = None
8 | max_new_tokens: Optional[int] = None
9 | repetition_penalty: Optional[float] = None
10 | temperature: Optional[float] = None
11 | top_k: Optional[float] = None
12 | top_p: Optional[float] = None
13 |
14 |
15 | class Query(BaseModel):
16 | query: str
17 | context: Optional[str] = None
18 | parameters: Optional[QueryParameters] = None
19 |
20 |
21 | class ChatMessage(BaseModel):
22 | role: Literal["user", "assistant", "function", "system"]
23 | content: str
24 |
25 |
26 | class ChatCompletion(BaseModel):
27 | model: str
28 | messages: List[ChatMessage]
29 |
--------------------------------------------------------------------------------
/src/schemas/training.py:
--------------------------------------------------------------------------------
1 | import yaml
2 | from pydantic import BaseModel
3 | from typing import Optional
4 | from src.schemas.deployment import Destination
5 | from src.yaml import loader, dumper
6 |
7 |
8 | class Hyperparameters(BaseModel):
9 | epochs: Optional[int] = None
10 | per_device_train_batch_size: Optional[int] = None
11 | learning_rate: Optional[float] = None
12 |
13 |
14 | class Training(BaseModel):
15 | destination: Destination
16 | instance_type: str
17 | instance_count: int
18 | training_input_path: str
19 | output_path: Optional[str] = None
20 | hyperparameters: Optional[Hyperparameters] = None
21 |
22 |
23 | def training_representer(dumper: yaml.SafeDumper, training: Training) -> yaml.nodes.MappingNode:
24 | return dumper.represent_mapping("!Training", {
25 | "destination": training.destination.value,
26 | "instance_type": training.instance_type,
27 | "instance_count": training.instance_count,
28 | "output_path": training.output_path,
29 | "training_input_path": training.training_input_path,
30 | "hyperparameters": training.hyperparameters,
31 | })
32 |
33 |
34 | def training_constructor(loader: yaml.SafeLoader, node: yaml.nodes.MappingNode) -> Training:
35 | return Training(**loader.construct_mapping(node))
36 |
37 |
38 | def hyperparameters_constructor(loader: yaml.SafeLoader, node: yaml.nodes.MappingNode) -> Hyperparameters:
39 | return Hyperparameters(**loader.construct_mapping(node))
40 |
41 |
42 | dumper.add_representer(Training, training_representer)
43 | loader.add_constructor(u'!Training', training_constructor)
44 | loader.add_constructor(u'!Hyperparameters', hyperparameters_constructor)
45 |
--------------------------------------------------------------------------------
/src/session.py:
--------------------------------------------------------------------------------
1 | import boto3
2 | import sagemaker
3 |
4 | session = boto3.session.Session()
5 | sagemaker_session = sagemaker.session.Session(boto_session=session)
6 |
--------------------------------------------------------------------------------
/src/utils/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/openfoundry-ai/model_manager/34f9ffc3e1962aa87ca088203535d8d0506b4e7b/src/utils/__init__.py
--------------------------------------------------------------------------------
/src/utils/aws_utils.py:
--------------------------------------------------------------------------------
1 | def construct_s3_uri(bucket: str, prefix: str) -> str:
2 | return f"s3://{bucket}/{prefix}"
3 |
4 |
5 | def is_s3_uri(path: str) -> bool:
6 | return path.startswith("s3://")
7 |
--------------------------------------------------------------------------------
/src/utils/format.py:
--------------------------------------------------------------------------------
1 | from typing import Dict
2 |
3 |
4 | def to_camel_case(snake_str: str) -> str:
5 | return "".join(x.capitalize() for x in snake_str.lower().split("_"))
6 |
7 |
8 | def to_lower_camel_case(snake_str: str) -> str:
9 | # We capitalize the first letter of each component except the first one
10 | # with the 'capitalize' method and join them together.
11 | camel_string = to_camel_case(snake_str)
12 | return lower_first(camel_string)
13 |
14 |
15 | def lower_first(string):
16 | return string[0].lower() + string[1:]
17 |
18 |
19 | def format_sagemaker_endpoint(output: Dict[str, str]) -> Dict[str, str]:
20 | return {lower_first(key): val for key, val in output.items()}
21 |
22 |
23 | def format_python_dict(output: Dict[str, str]):
24 | return {to_lower_camel_case(key): val for key, val in output.items()}
25 |
--------------------------------------------------------------------------------
/src/utils/model_utils.py:
--------------------------------------------------------------------------------
1 | import datetime
2 | from difflib import SequenceMatcher
3 | from dotenv import dotenv_values
4 | from huggingface_hub import HfApi
5 | from src.utils.rich_utils import print_error
6 | from src.sagemaker import SagemakerTask
7 | from src.schemas.deployment import Deployment
8 | from src.schemas.model import Model, ModelSource
9 | from src.schemas.query import Query
10 | from src.session import sagemaker_session
11 | from typing import Dict, Tuple, Optional
12 | HUGGING_FACE_HUB_TOKEN = dotenv_values(".env").get("HUGGING_FACE_HUB_KEY")
13 |
14 |
15 | def get_unique_endpoint_name(model_id: str, endpoint_name: str = None):
16 | dt_string = datetime.datetime.now().strftime("%Y%m%d%H%M")
17 |
18 | if not endpoint_name:
19 | # Endpoint name must be < 63 characters
20 | model_string = model_id.replace(
21 | "/", "--").replace("_", "-").replace(".", "")[:50]
22 | return f"{model_string}-{dt_string}"
23 | else:
24 | return f"{endpoint_name[:50]}-{dt_string}"
25 |
26 |
27 | def is_sagemaker_model(endpoint_name: str, config: Optional[Tuple[Deployment, Model]] = None) -> bool:
28 | if config is not None:
29 | _, model = config
30 | task = get_sagemaker_model_and_task(model.id)['task']
31 |
32 | # check task for custom sagemaker models
33 | return model.source == ModelSource.Sagemaker or task in SagemakerTask.list()
34 |
35 | # fallback
36 | return endpoint_name.find("--") == -1
37 |
38 |
39 | def is_custom_model(endpoint_name: str) -> bool:
40 | return endpoint_name.startswith('custom')
41 |
42 |
43 | def get_sagemaker_model_and_task(endpoint_or_model_name: str):
44 | components = endpoint_or_model_name.split('-')
45 | framework, task = components[:2]
46 | model_id = '-'.join(components[:-1])
47 | return {
48 | 'model_id': model_id,
49 | 'task': task,
50 | }
51 |
52 |
53 | def get_hugging_face_pipeline_task(model_name: str):
54 | hf_api = HfApi()
55 | try:
56 | model_info = hf_api.model_info(
57 | model_name, token=HUGGING_FACE_HUB_TOKEN)
58 | task = model_info.pipeline_tag
59 | except Exception:
60 | print_error("Model not found, please try another.")
61 | return None
62 |
63 | return task
64 |
65 |
66 | def get_model_and_task(endpoint_or_model_name: str, config: Optional[Tuple[Deployment, Model]] = None) -> Dict[str, str]:
67 | if config is not None:
68 | _, model = config
69 | return {
70 | 'model_id': model.id,
71 | 'task': model.task
72 | }
73 |
74 | if (is_sagemaker_model(endpoint_or_model_name)):
75 | return get_sagemaker_model_and_task(endpoint_or_model_name)
76 | else:
77 | model_id = get_model_name_from_hugging_face_endpoint(
78 | endpoint_or_model_name)
79 | task = get_hugging_face_pipeline_task(model_id)
80 | return {
81 | 'model_id': model_id,
82 | 'task': task
83 | }
84 |
85 |
86 | def get_model_name_from_hugging_face_endpoint(endpoint_name: str):
87 | endpoint_name = endpoint_name.removeprefix("custom-")
88 | endpoint_name = endpoint_name.replace("--", "/")
89 | author, rest = endpoint_name.split("/")
90 |
91 | # remove datetime
92 | split = rest.split('-')
93 | fuzzy_model_name = '-'.join(split[:-1])
94 |
95 | # get first token
96 | search_term = fuzzy_model_name.split('-')[0]
97 |
98 | hf_api = HfApi()
99 | results = hf_api.list_models(search=search_term, author=author)
100 |
101 | # find results that closest match our fuzzy model name
102 | results_to_diff = {}
103 | for result in results:
104 | results_to_diff[result.id] = SequenceMatcher(
105 | None, result.id, f"{author}/{fuzzy_model_name}").ratio()
106 |
107 | return max(results_to_diff, key=results_to_diff.get)
108 |
109 |
110 | def get_text_generation_hyperpameters(config: Optional[Tuple[Deployment, Model]], query: Query = None):
111 | if query.parameters is not None:
112 | return query.parameters.model_dump()
113 |
114 | if config is not None and config[1].predict is not None:
115 | return config[1].predict
116 |
117 | # Defaults
118 | return {
119 | "max_new_tokens": 250,
120 | "top_p": 0.9,
121 | "temperature": 0.9,
122 | }
123 |
--------------------------------------------------------------------------------
/src/utils/rich_utils.py:
--------------------------------------------------------------------------------
1 | from rich import print
2 |
3 |
4 | def print_model(text: str):
5 | print(f"[blue] {text}")
6 |
7 |
8 | def print_error(text: str):
9 | print(f"[red] {text}")
10 |
11 |
12 | def print_success(text: str):
13 | print(f"[green] {text}")
14 |
--------------------------------------------------------------------------------
/src/yaml.py:
--------------------------------------------------------------------------------
1 | import yaml
2 |
3 | loader = yaml.SafeLoader
4 | dumper = yaml.SafeDumper
5 |
--------------------------------------------------------------------------------
/tailwind.config.js:
--------------------------------------------------------------------------------
1 | /** @type {import('tailwindcss').Config} */
2 | module.exports = {
3 | content: [
4 | "./app/**/*.{js,ts,jsx,tsx,mdx}",
5 | "./components/**/*.{js,ts,jsx,tsx,mdx}",
6 | ],
7 | theme: {
8 | extend: {
9 | colors: {
10 | background: "hsl(var(--background))",
11 | foreground: "hsl(var(--foreground))",
12 | btn: {
13 | background: "hsl(var(--btn-background))",
14 | "background-hover": "hsl(var(--btn-background-hover))",
15 | },
16 | },
17 | },
18 | },
19 | plugins: [],
20 | };
21 |
--------------------------------------------------------------------------------
/tests/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/openfoundry-ai/model_manager/34f9ffc3e1962aa87ca088203535d8d0506b4e7b/tests/__init__.py
--------------------------------------------------------------------------------
/tests/test_openai_proxy.py:
--------------------------------------------------------------------------------
1 | import httpx
2 | from openai import OpenAI, DefaultHttpxClient
3 |
4 |
5 | def test_openai_proxy():
6 | client = OpenAI(
7 | api_key="test-key",
8 | base_url="http://127.0.0.1:8000/",
9 | http_client=DefaultHttpxClient(
10 | proxy="http://127.0.0.1:8000/",
11 | transport=httpx.HTTPTransport(local_address="0.0.0.0"),
12 | )
13 | )
14 | chat_completion = client.chat.completions.create(
15 | model="meta-llama/Meta-Llama-3-8B",
16 | messages=[
17 | {
18 | "role": "system",
19 | "content": "You are a helpful assistant.",
20 | },
21 | {
22 | "role": "user",
23 | "content": "Say this is a test",
24 | },
25 | ],
26 | )
27 | print(chat_completion)
28 |
--------------------------------------------------------------------------------
/tsconfig.json:
--------------------------------------------------------------------------------
1 | {
2 | "compilerOptions": {
3 | "target": "es5",
4 | "lib": ["dom", "dom.iterable", "esnext"],
5 | "allowJs": true,
6 | "skipLibCheck": true,
7 | "strict": true,
8 | "forceConsistentCasingInFileNames": true,
9 | "noEmit": true,
10 | "esModuleInterop": true,
11 | "module": "esnext",
12 | "moduleResolution": "node",
13 | "resolveJsonModule": true,
14 | "isolatedModules": true,
15 | "jsx": "preserve",
16 | "incremental": true,
17 | "plugins": [
18 | {
19 | "name": "next"
20 | }
21 | ],
22 | "paths": {
23 | "@/*": ["./*"]
24 | }
25 | },
26 | "include": ["next-env.d.ts", "**/*.ts", "**/*.tsx", ".next/types/**/*.ts"],
27 | "exclude": ["node_modules"]
28 | }
29 |
--------------------------------------------------------------------------------
/utils/supabase/client.ts:
--------------------------------------------------------------------------------
1 | import { createBrowserClient } from "@supabase/ssr";
2 |
3 | export const createClient = () =>
4 | createBrowserClient(
5 | process.env.NEXT_PUBLIC_SUPABASE_URL!,
6 | process.env.NEXT_PUBLIC_SUPABASE_ANON_KEY!,
7 | );
8 |
--------------------------------------------------------------------------------
/utils/supabase/middleware.ts:
--------------------------------------------------------------------------------
1 | import { createServerClient, type CookieOptions } from "@supabase/ssr";
2 | import { type NextRequest, NextResponse } from "next/server";
3 |
4 | export const updateSession = async (request: NextRequest) => {
5 | // This `try/catch` block is only here for the interactive tutorial.
6 | // Feel free to remove once you have Supabase connected.
7 | try {
8 | // Create an unmodified response
9 | let response = NextResponse.next({
10 | request: {
11 | headers: request.headers,
12 | },
13 | });
14 |
15 | const supabase = createServerClient(
16 | process.env.NEXT_PUBLIC_SUPABASE_URL!,
17 | process.env.NEXT_PUBLIC_SUPABASE_ANON_KEY!,
18 | {
19 | cookies: {
20 | get(name: string) {
21 | return request.cookies.get(name)?.value;
22 | },
23 | set(name: string, value: string, options: CookieOptions) {
24 | // If the cookie is updated, update the cookies for the request and response
25 | request.cookies.set({
26 | name,
27 | value,
28 | ...options,
29 | });
30 | response = NextResponse.next({
31 | request: {
32 | headers: request.headers,
33 | },
34 | });
35 | response.cookies.set({
36 | name,
37 | value,
38 | ...options,
39 | });
40 | },
41 | remove(name: string, options: CookieOptions) {
42 | // If the cookie is removed, update the cookies for the request and response
43 | request.cookies.set({
44 | name,
45 | value: "",
46 | ...options,
47 | });
48 | response = NextResponse.next({
49 | request: {
50 | headers: request.headers,
51 | },
52 | });
53 | response.cookies.set({
54 | name,
55 | value: "",
56 | ...options,
57 | });
58 | },
59 | },
60 | },
61 | );
62 |
63 | // This will refresh session if expired - required for Server Components
64 | // https://supabase.com/docs/guides/auth/server-side/nextjs
65 | await supabase.auth.getUser();
66 |
67 | return response;
68 | } catch (e) {
69 | // If you are here, a Supabase client could not be created!
70 | // This is likely because you have not set up environment variables.
71 | // Check out http://localhost:3000 for Next Steps.
72 | return NextResponse.next({
73 | request: {
74 | headers: request.headers,
75 | },
76 | });
77 | }
78 | };
79 |
--------------------------------------------------------------------------------
/utils/supabase/server.ts:
--------------------------------------------------------------------------------
1 | import { createServerClient, type CookieOptions } from "@supabase/ssr";
2 | import { cookies } from "next/headers";
3 |
4 | export const createClient = () => {
5 | const cookieStore = cookies();
6 |
7 | return createServerClient(
8 | process.env.NEXT_PUBLIC_SUPABASE_URL!,
9 | process.env.NEXT_PUBLIC_SUPABASE_ANON_KEY!,
10 | {
11 | cookies: {
12 | get(name: string) {
13 | return cookieStore.get(name)?.value;
14 | },
15 | set(name: string, value: string, options: CookieOptions) {
16 | try {
17 | cookieStore.set({ name, value, ...options });
18 | } catch (error) {
19 | // The `set` method was called from a Server Component.
20 | // This can be ignored if you have middleware refreshing
21 | // user sessions.
22 | }
23 | },
24 | remove(name: string, options: CookieOptions) {
25 | try {
26 | cookieStore.set({ name, value: "", ...options });
27 | } catch (error) {
28 | // The `delete` method was called from a Server Component.
29 | // This can be ignored if you have middleware refreshing
30 | // user sessions.
31 | }
32 | },
33 | },
34 | },
35 | );
36 | };
37 |
--------------------------------------------------------------------------------