wehub-resource-sync
113 lines
3.1 KiB
Python
113 lines
3.1 KiB
Python
"""
|
|
Usage
|
|
-----
|
|
databricks secrets create-scope <scope>
|
|
databricks secrets put-secret <scope> openai-api-key --string-value $OPENAI_API_KEY
|
|
python examples/deployments/databricks.py --secret <scope>/openai-api-key
|
|
-----
|
|
"""
|
|
|
|
import argparse
|
|
import uuid
|
|
|
|
from mlflow.deployments import get_deploy_client
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--secret", type=str, help="Secret (e.g. secrets/scope/key)")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
client = get_deploy_client("databricks")
|
|
name = f"test-endpoint-{uuid.uuid4()}"
|
|
client.create_endpoint(
|
|
name=name,
|
|
config={
|
|
"served_entities": [
|
|
{
|
|
"name": "test",
|
|
"external_model": {
|
|
"name": "gpt-4",
|
|
"provider": "openai",
|
|
"task": "llm/v1/chat",
|
|
"openai_config": {
|
|
"openai_api_key": "{{" + args.secret + "}}",
|
|
},
|
|
},
|
|
}
|
|
],
|
|
"tags": [
|
|
{
|
|
"key": "foo",
|
|
"value": "bar",
|
|
}
|
|
],
|
|
"rate_limits": [
|
|
{
|
|
"key": "user",
|
|
"renewal_period": "minute",
|
|
"calls": 5,
|
|
}
|
|
],
|
|
},
|
|
)
|
|
try:
|
|
# Update served_entities
|
|
print(
|
|
client.update_endpoint(
|
|
endpoint=name,
|
|
config={
|
|
"served_entities": [
|
|
{
|
|
"name": "test",
|
|
"external_model": {
|
|
"name": "gpt-4",
|
|
"provider": "openai",
|
|
"task": "llm/v1/chat",
|
|
"openai_config": {
|
|
"openai_api_key": "{{" + args.secret + "}}",
|
|
},
|
|
},
|
|
}
|
|
],
|
|
},
|
|
)
|
|
)
|
|
# Update rate_limits
|
|
print(
|
|
client.update_endpoint(
|
|
endpoint=name,
|
|
config={
|
|
"rate_limits": [
|
|
{
|
|
"key": "user",
|
|
"renewal_period": "minute",
|
|
"calls": 10,
|
|
}
|
|
],
|
|
},
|
|
)
|
|
)
|
|
print(client.list_endpoints()[:3])
|
|
print(client.get_endpoint(endpoint=name))
|
|
print(
|
|
client.predict(
|
|
endpoint=name,
|
|
inputs={
|
|
"messages": [
|
|
{"role": "user", "content": "Hello!"},
|
|
],
|
|
"max_tokens": 128,
|
|
},
|
|
),
|
|
)
|
|
finally:
|
|
client.delete_endpoint(endpoint=name)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|