diff --git a/pyproject.toml b/pyproject.toml index bfa33ae37e0d..2612023fd2d9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -138,9 +138,6 @@ test = [ "weave[trace_server_tests]", ] -[project.scripts] -weave = "weave.trace.cli:cli" - [project.urls] Company = "https://wandb.com" Documentation = "https://docs.wandb.com/" diff --git a/tests/trace/test_cli.py b/tests/trace/test_cli.py deleted file mode 100644 index 9710611e85e1..000000000000 --- a/tests/trace/test_cli.py +++ /dev/null @@ -1,11 +0,0 @@ -from click.testing import CliRunner - -from weave.trace.cli import cli -from weave.version import VERSION - - -def test_version(): - runner = CliRunner() - result = runner.invoke(cli, ["--version"]) - assert result.exit_code == 0 - assert result.output == f"cli, version {VERSION}\n" diff --git a/weave/deploy/Dockerfile b/weave/deploy/Dockerfile deleted file mode 100644 index 832d26464250..000000000000 --- a/weave/deploy/Dockerfile +++ /dev/null @@ -1,11 +0,0 @@ -FROM $BASE_IMAGE - -ENV PYTHONUNBUFFERED 1 -WORKDIR /app - -COPY requirements/* . -RUN pip install --no-cache-dir -r requirements.txt -ENV PROJECT_NAME $PROJECT_NAME - -EXPOSE 8080 -CMD ["weave", "serve", "$MODEL_REF", "--port=8080", "--method=$MODEL_METHOD", "--auth-entity=$AUTH_ENTITY"] diff --git a/weave/deploy/__init__.py b/weave/deploy/__init__.py deleted file mode 100644 index 8b137891791f..000000000000 --- a/weave/deploy/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/weave/deploy/gcp/__init__.py b/weave/deploy/gcp/__init__.py deleted file mode 100644 index b73b6aa9b3f7..000000000000 --- a/weave/deploy/gcp/__init__.py +++ /dev/null @@ -1,345 +0,0 @@ -import os -import shutil -import string -import subprocess -import sys -import tempfile -import typing -from pathlib import Path - -from weave import __version__ -from weave.deploy.util import execute, safe_name -from weave.trace import env -from weave.trace.refs import ObjectRef, parse_uri - - -def generate_dockerfile( - model_ref: str, - model_method: typing.Optional[str] = None, - project_name: typing.Optional[str] = None, - auth_entity: typing.Optional[str] = None, - base_image: typing.Optional[str] = "python:3.11", -) -> str: - """Generates a Dockerfile to run the weave op""" - if project_name is None: - ref_uri = parse_uri(model_ref) - if not isinstance(ref_uri, ObjectRef): - raise ValueError(f"Expected a ObjectRef artifact ref, got {type(ref_uri)}") - project_name = ref_uri.project - src = Path(__file__).parent.parent / "Dockerfile" - template = string.Template(src.read_text()) - - return template.substitute( - { - "PROJECT_NAME": project_name, - "BASE_IMAGE": base_image, - "MODEL_REF": model_ref, - "MODEL_METHOD": model_method or "", - "AUTH_ENTITY": auth_entity or "", - } - ) - - -def generate_requirements_txt(model_ref: str, dir: str, dev: bool = False) -> str: - """Generate a requirements.txt file.""" - cwd = Path(os.getcwd()) - if dev and (cwd / "build_dist.py").exists(): - print("Building weave for development...") - env_dict = os.environ.copy() - env_dict.update({"WEAVE_SKIP_BUILD": "1"}) - execute( - [sys.executable, str(cwd / "build_dist.py")], env=env_dict, capture=False - ) - wheel = f"weave-{__version__}-py3-none-any.whl" - execute(["cp", str(cwd / "dist" / wheel), dir], capture=False) - weave = f"/app/{wheel}" - else: - weave = "weave @ git+https://github.com/wandb/weave@master" - # TODO: add any additional reqs the op needs - - # We're requiring faiss-cpu for now here, to get Hooman Slackbot deploy - # working. But this is not right, objects and ops should have their own - # requirements that we compile together here. - # TODO: Fix - return f""" -uvicorn[standard] -fastapi -faiss-cpu -{weave} -""" - - -def gcloud( - args: list[str], - timeout: typing.Optional[float] = None, - input: typing.Optional[str] = None, - capture: bool = True, -) -> typing.Any: - gcloud_absolute_path = shutil.which("gcloud") - if gcloud_absolute_path is None: - raise ValueError( - "gcloud command required: https://cloud.google.com/sdk/docs/install" - ) - if os.getenv("DEBUG") == "true": - print(f"Running gcloud {' '.join(args)}") - return execute( - [gcloud_absolute_path] + args, timeout=timeout, capture=capture, input=input - ) - - -def enforce_login() -> None: - """Ensure the user is logged in to gcloud.""" - try: - auth = gcloud(["auth", "print-access-token", "--format=json"], timeout=3) - if auth.get("token") is None: - raise ValueError() - except (subprocess.TimeoutExpired, ValueError): - raise ValueError("Not logged in to gcloud. Please run `gcloud auth login`.") - - -def compile( - model_ref: str, - model_method: typing.Optional[str] = None, - wandb_project: typing.Optional[str] = None, - auth_entity: typing.Optional[str] = None, - base_image: typing.Optional[str] = None, - dev: bool = False, -) -> str: - """Compile the weave application.""" - dir = tempfile.mkdtemp() - reqs = os.path.join(dir, "requirements") - os.mkdir(reqs) - with open(os.path.join(reqs, "requirements.txt"), "w") as f: - f.write(generate_requirements_txt(model_ref, reqs, dev)) - with open(os.path.join(dir, "Dockerfile"), "w") as f: - f.write( - generate_dockerfile( - model_ref, model_method, wandb_project, auth_entity, base_image - ) - ) - return dir - - -def ensure_service_account( - name: str = "weave-default", project: typing.Optional[str] = None -) -> str: - """Ensure the user has a service account.""" - if len(name) < 6 or len(name) > 30: - raise ValueError("Service account name must be between 6 and 30 characters.") - project = project or gcloud(["config", "get", "project", "--format=json"]) - account = gcloud(["auth", "list", "--filter=status:ACTIVE", "--format=json"])[0][ - "account" - ] - sa = f"{name}@{project}.iam.gserviceaccount.com" - exists = gcloud( - ["iam", "service-accounts", "list", f"--filter=email={sa}", "--format=json"] - ) - if len(exists) == 0: - print(f"Creating service account {name}...") - display_name = ( - " ".join([n.capitalize() for n in name.split("-")]) + " Service Account" - ) - gcloud( - [ - "iam", - "service-accounts", - "create", - name, - f"--display-name={display_name}", - f"--project={project}", - "--format=json", - ] - ) - gcloud( - [ - "iam", - "service-accounts", - "add-iam-policy-binding", - f"{sa}", - f"--project={project}", - f"--member=user:{account}", - "--role=roles/iam.serviceAccountUser", - "--format=json", - ] - ) - print( - "To grant additional permissions, run add-iam-policy-binding on the resource:" - ) - print( - " gcloud storage buckets add-iam-policy-binding gs://BUCKET --member=serviceAccount:{sa} --role=ROLE" - ) - else: - print(f"Using service account: {sa}") - return sa - - -def ensure_secret( - name: str, value: str, service_account: str, project: typing.Optional[str] = None -) -> None: - """Ensure a secret exists and is accessbile by the service account.""" - project = project or gcloud(["config", "get", "project", "--format=json"]) - exists = gcloud( - [ - "secrets", - "list", - rf"--filter=name~^.*\/{name}$", - f"--project={project}", - "--format=json", - ] - ) - if len(exists) == 0: - print(f"Creating secret {name} and granting access to {service_account}...") - gcloud( - [ - "secrets", - "create", - name, - f"--project={project}", - "--replication-policy=automatic", - "--format=json", - ] - ) - # To support changing service accounts, we always add secretAccessor - gcloud( - [ - "secrets", - "add-iam-policy-binding", - name, - f"--project={project}", - "--format=json", - f"--member=serviceAccount:{service_account}", - "--role=roles/secretmanager.secretAccessor", - ] - ) - gcloud( - [ - "secrets", - "versions", - "add", - name, - f"--project={project}", - "--format=json", - "--data-file=-", - ], - input=value, - ) - - -# This is a sketch or the commands needed to downscope permissions and use secrets -def deploy( - model_ref: str, - model_method: typing.Optional[str] = None, - wandb_project: typing.Optional[str] = None, - gcp_project: typing.Optional[str] = None, - region: typing.Optional[str] = None, - service_account: typing.Optional[str] = None, - auth_entity: typing.Optional[str] = None, - base_image: typing.Optional[str] = "python:3.11", - memory: typing.Optional[str] = "1000Mi", -) -> None: - """Deploy the weave application.""" - enforce_login() - if region is None: - region = gcloud(["config", "get", "functions/region", "--format=json"]) - if region == []: - raise ValueError( - "No default region set. Run `gcloud config set functions/region ` or set the region argument." - ) - if service_account is None: - try: - service_account = ensure_service_account(project=gcp_project) - except ValueError: - print( - "WARNING: No service account specified. Using the compute engine default service account..." - ) - dir = compile(model_ref, model_method, wandb_project, auth_entity, base_image) - ref = parse_uri(model_ref) - if not isinstance(ref, ObjectRef): - raise TypeError(f"Expected a weave object uri, got {type(ref)}") - name = safe_name(f"{ref.project}-{ref.name}") - project = wandb_project or ref.project - key = env.weave_wandb_api_key() - secrets = { - "WANDB_API_KEY": key, - } - if os.getenv("OPENAI_API_KEY"): - secrets["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY") - args = [ - "run", - "deploy", - name, - f"--region={region}", - f"--memory={memory}", - f"--set-env-vars=PROJECT_NAME={project}", - f"--source={dir}", - "--allow-unauthenticated", - ] - sec_or_env = "" - if service_account is not None: - args.append(f"--service-account={service_account}") - sec_or_env = "--set-secrets=" - for k, v in secrets.items(): - if v is not None: - ensure_secret(k, v, service_account, gcp_project) - sec_or_env += f"{k}={k}:latest," - else: - sec_or_env = "--set-env-vars=" - for k, v in secrets.items(): - sec_or_env += f"{k}={v}," - # trim the trailing comma - sec_or_env = sec_or_env[:-1] - args.append(sec_or_env) - if gcp_project is not None: - args.append(f"--project={gcp_project}") - gcloud(args, capture=False) - shutil.rmtree(dir) - - -def develop( - model_ref: str, - model_method: typing.Optional[str] = None, - base_image: typing.Optional[str] = "python:3.11", - auth_entity: typing.Optional[str] = None, -) -> None: - dir = compile( - model_ref, - model_method=model_method, - base_image=base_image, - auth_entity=auth_entity, - dev=True, - ) - model_uri = parse_uri(model_ref) - if not isinstance(model_uri, ObjectRef): - raise TypeError(f"Expected a weave object uri, got {type(model_uri)}") - name = safe_name(model_uri.name) - docker = shutil.which("docker") - if docker is None: - raise ValueError("docker command required: https://docs.docker.com/get-docker/") - print("Building container from: ", dir) - execute( - [docker, "buildx", "build", "-t", name, "--load", "."], cwd=dir, capture=False - ) - env_api_key = env.weave_wandb_api_key() - if env_api_key is None: - raise ValueError("WANDB_API_KEY environment variable required") - env_dict = {"WANDB_API_KEY": env_api_key} - env_dict.update(os.environ.copy()) - print("Running container at http://localhost:8080") - execute( - [ - docker, - "run", - "-p", - "8080:8080", - "-e", - "WANDB_API_KEY", - "-e", - "OPENAI_API_KEY", - name, - ], - env=env_dict, - capture=False, - ) - if os.getenv("DEBUG") == None: - print("Cleaning up...") - shutil.rmtree(dir) diff --git a/weave/deploy/modal/__init__.py b/weave/deploy/modal/__init__.py deleted file mode 100644 index 8a5d35f35bf5..000000000000 --- a/weave/deploy/modal/__init__.py +++ /dev/null @@ -1,124 +0,0 @@ -import os -import string -import tempfile -import time -import typing -from pathlib import Path - -from weave_query import artifact_wandb as artifact_wandb # type: ignore -from weave_query import environment # type: ignore - -from weave.trace.refs import ObjectRef, parse_uri - -try: - from modal.cli.import_refs import import_stub - from modal.config import config - from modal.runner import deploy_stub - from modal.serving import serve_stub -except ImportError: - raise ImportError( - "modal must be installed and configured: \n pip install weave[modal]\n modal setup" - ) - - -def compile( - model_ref: str, - project: str, - reqs: list[str], - auth_entity: typing.Optional[str] = None, - secrets: typing.Optional[dict[str, str]] = None, -) -> Path: - """Generates a modal py file and secret env vars to run the weave op""" - dir = Path(tempfile.mkdtemp()) - with open(Path(__file__).parent / "stub.py") as f: - template = string.Template(f.read()) - src = template.substitute( - { - "REQUIREMENTS": '", "'.join(reqs), - "MODEL_REF": model_ref, - "PROJECT_NAME": project, - "AUTH_ENTITY": auth_entity or "", - } - ) - - with open(dir / "weave_model.py", "w") as f: - f.write(src) - with open(dir / ".env", "w") as f: - if secrets is not None: - for k, v in secrets.items(): - f.write(f"{k}={v}\n") - return dir - - -def generate_modal_stub( - model_ref: str, - project_name: typing.Optional[str] = None, - reqs: typing.Optional[list[str]] = None, - auth_entity: typing.Optional[str] = None, - secrets: typing.Optional[dict[str, str]] = None, -) -> str: - """Generates a modal py file to run the weave op""" - parsed_ref = parse_uri(model_ref) - if project_name is None: - if not isinstance(parsed_ref, ObjectRef): - raise ValueError(f"Expected a weave object uri, got {type(parsed_ref)}") - project_name = parsed_ref.project - - project = project_name or os.getenv("PROJECT_NAME") - if project is None: - raise ValueError( - "project must be specified from command line or via the PROJECT_NAME env var" - ) - reqs = reqs or [] - reqs.append("weave @ git+https://github.com/wandb/weave@weaveflow") - reqs.append("fastapi>=0.104.0") - return str( - compile(model_ref, project, reqs, secrets=secrets, auth_entity=auth_entity) - / "weave_model.py" - ) - - -def extract_secrets(model_ref: str) -> dict[str, str]: - # TODO: get secrets from the weave op - key = environment.weave_wandb_api_key() - if key is None: - secrets = {} - else: - secrets = { - "WANDB_API_KEY": key, - } - openai_api_key = os.getenv("OPENAI_API_KEY") - if openai_api_key: - secrets["OPENAI_API_KEY"] = openai_api_key - return secrets - - -def deploy( - model_ref: str, - wandb_project: typing.Optional[str] = None, - auth_entity: typing.Optional[str] = None, -) -> None: - """Deploy a model to the modal labs cloud.""" - ref = generate_modal_stub( - model_ref, - wandb_project, - secrets=extract_secrets(model_ref), - auth_entity=auth_entity, - ) - stub = import_stub(ref) - deploy_stub(stub, name=stub.name, environment_name=config.get("environment")) - - -def develop(model_ref: str, auth_entity: typing.Optional[str] = None) -> None: - """Run a model for testing.""" - ref = generate_modal_stub( - model_ref, secrets=extract_secrets(model_ref), auth_entity=auth_entity - ) - print(f"Serving live code from: {ref}") - stub = import_stub(ref) - timeout = 1800 - with serve_stub(stub, ref, environment_name=config.get("environment")): - while timeout > 0: - t = min(timeout, 3600) - time.sleep(t) - timeout -= t diff --git a/weave/deploy/modal/stub.py b/weave/deploy/modal/stub.py deleted file mode 100644 index 9975ac343a00..000000000000 --- a/weave/deploy/modal/stub.py +++ /dev/null @@ -1,40 +0,0 @@ -import os - -from fastapi import FastAPI -from modal import Image, Secret, Stub, asgi_app -from weave_query.uris import WeaveURI - -from weave.deploy.util import safe_name -from weave.trace.refs import ObjectRef, parse_uri - -image = ( - Image.debian_slim() - .apt_install("git") - .pip_install(["$REQUIREMENTS"]) - .env( - { - "MODEL_REF": "$MODEL_REF", - "PROJECT_NAME": "$PROJECT_NAME", - "AUTH_ENTITY": "$AUTH_ENTITY", - } - ) -) -stub = Stub("$PROJECT_NAME") -uri = WeaveURI.parse("$MODEL_REF") - - -@stub.function(image=image, secret=Secret.from_dotenv(__file__)) -@asgi_app(label=safe_name(uri.name)) -def fastapi_app() -> FastAPI: - from weave.trace import api - from weave.trace.serve_fastapi import object_method_app - - uri_ref = parse_uri(os.environ["MODEL_REF"]) - if not isinstance(uri_ref, ObjectRef): - raise TypeError(f"Expected a weave object uri, got {type(uri_ref)}") - app = object_method_app(uri_ref, auth_entity=os.environ.get("AUTH_ENTITY")) - - api.init(os.environ["PROJECT_NAME"]) - # TODO: hookup / provide more control over attributes - # with api.attributes({"env": env}): - return app diff --git a/weave/deploy/util.py b/weave/deploy/util.py deleted file mode 100644 index 3e862fec120a..000000000000 --- a/weave/deploy/util.py +++ /dev/null @@ -1,48 +0,0 @@ -import json -import os -import re -import subprocess -import sys -import typing - - -def execute( - args: list[str], - timeout: typing.Optional[float] = None, - cwd: typing.Optional[str] = None, - env: typing.Optional[dict[str, str]] = None, - input: typing.Optional[str] = None, - capture: bool = True, -) -> typing.Any: - process = subprocess.Popen( - args, - stdout=subprocess.PIPE if capture else sys.stdout, - stderr=subprocess.PIPE if capture else sys.stderr, - stdin=subprocess.PIPE if capture else sys.stdin, - universal_newlines=True, - env=env or os.environ.copy(), - cwd=cwd, - ) - out, err = process.communicate(timeout=timeout, input=input) - if process.returncode != 0: - raise ValueError(f"Command failed: {err or ''}") - - if not capture: - return None - - try: - return json.loads(out) - except json.JSONDecodeError: - raise ValueError(f"Failed to parse JSON from command: {out}") - - -def safe_name(name: str) -> str: - """The name must use only lowercase alphanumeric characters and dashes, - cannot begin or end with a dash, and cannot be longer than 63 characters.""" - fixed_name = re.sub(r"[^a-z0-9-]", "-", name.lower()).strip("-") - if len(fixed_name) == 0: - return "weave-op" - elif len(fixed_name) > 63: - return fixed_name[:63] - else: - return fixed_name diff --git a/weave/trace/api.py b/weave/trace/api.py index cbbce9adb11c..13901662d280 100644 --- a/weave/trace/api.py +++ b/weave/trace/api.py @@ -3,16 +3,13 @@ from __future__ import annotations import contextlib -import os -import threading -import time from collections.abc import Iterator from typing import Any # TODO: type_handlers is imported here to trigger registration of the image serializer. # There is probably a better place for this, but including here for now to get the fix in. from weave import type_handlers # noqa: F401 -from weave.trace import urls, util, weave_client, weave_init +from weave.trace import urls, weave_client, weave_init from weave.trace.autopatch import AutopatchSettings from weave.trace.constants import TRACE_OBJECT_EMOJI from weave.trace.context import call_context @@ -231,53 +228,6 @@ def attributes(attributes: dict[str, Any]) -> Iterator: call_context.call_attributes.reset(token) -def serve( - model_ref: ObjectRef, - method_name: str | None = None, - auth_entity: str | None = None, - port: int = 9996, - thread: bool = False, -) -> str: - import uvicorn - - from weave.trace.serve_fastapi import object_method_app - from weave.wandb_interface import wandb_api - - client = weave_client_context.require_weave_client() - # if not isinstance( - # client, _graph_client_wandb_art_st.GraphClientWandbArtStreamTable - # ): - # raise ValueError("serve currently only supports wandb client") - - print(f"Serving {model_ref}") - print(f"🥐 Server docs and playground at http://localhost:{port}/docs") - print() - os.environ["PROJECT_NAME"] = f"{client.entity}/{client.project}" - os.environ["MODEL_REF"] = str(model_ref) - - wandb_api_ctx = wandb_api.get_wandb_api_context() - app = object_method_app(model_ref, method_name=method_name, auth_entity=auth_entity) - trace_attrs = call_context.call_attributes.get() - - def run() -> None: - # This function doesn't return, because uvicorn.run does not return. - with wandb_api.wandb_api_context(wandb_api_ctx): - with attributes(trace_attrs): - uvicorn.run(app, host="0.0.0.0", port=port) - - if util.is_notebook(): - thread = True - if thread: - t = threading.Thread(target=run, daemon=True) - t.start() - time.sleep(1) - return "http://localhost:%d" % port - else: - # Run should never return - run() - raise ValueError("Should not reach here") - - def finish() -> None: """Stops logging to weave. @@ -303,7 +253,6 @@ def finish() -> None: "obj_ref", "output_of", "attributes", - "serve", "finish", "op", "Table", diff --git a/weave/trace/cli.py b/weave/trace/cli.py deleted file mode 100644 index cc327ecbd1f7..000000000000 --- a/weave/trace/cli.py +++ /dev/null @@ -1,141 +0,0 @@ -import os -import typing - -import click - -from weave import __version__ -from weave.deploy import gcp as google -from weave.trace import api -from weave.trace.refs import ObjectRef, parse_uri - -# TODO: does this work? -os.environ["PYTHONUNBUFFERED"] = "1" - - -@click.group() -@click.version_option(version=__version__) -def cli() -> None: - pass - - -@cli.command(help="Serve weave models.") # pyright: ignore[reportFunctionMemberAccess] -@click.argument("model_ref") -@click.option("--method", help="Method name to serve.") -@click.option("--project", help="W&B project name.") -@click.option("--env", default="development", help="Environment to tag the model with.") -@click.option( - "--auth-entity", help="Enforce authorization for this entity using W&B API keys." -) -@click.option("--port", default=9996, type=int) -def serve( - model_ref: str, - method: typing.Optional[str], - auth_entity: typing.Optional[str], - project: str, - env: str, - port: int, -) -> None: - parsed_ref = parse_uri(model_ref) - if not isinstance(parsed_ref, ObjectRef): - raise TypeError(f"Expected a weave artifact uri, got {parsed_ref}") - ref_project = parsed_ref.project - project_override = project or os.getenv("PROJECT_NAME") - if project_override: - print(f"Logging to project different from {ref_project}") - project = project_override - else: - project = ref_project - - api.init(project) - # TODO: provide more control over attributes - with api.attributes({"env": env}): - api.serve( - parsed_ref, method_name=method or None, auth_entity=auth_entity, port=port - ) - - -@cli.group(help="Deploy weave models.") # pyright: ignore[reportFunctionMemberAccess] -def deploy() -> None: - pass - - -@deploy.command(help="Deploy to GCP.") -@click.argument("model_ref") -@click.option("--method", help="Method name to serve.") -@click.option("--project", help="W&B project name.") -@click.option("--gcp-project", help="GCP project name.") -@click.option( - "--auth-entity", help="Enforce authorization for this entity using W&B API keys." -) -@click.option("--service-account", help="GCP service account.") -@click.option("--dev", is_flag=True, help="Run the function locally.") -def gcp( - model_ref: str, - method: str, - project: str, - gcp_project: str, - auth_entity: str, - service_account: str, - dev: bool = False, -) -> None: - if dev: - print(f"Developing model {model_ref}...") - google.develop(model_ref, model_method=method, auth_entity=auth_entity) - return - print(f"Deploying model {model_ref}...") - if auth_entity is None: - print( - "WARNING: No --auth-entity specified. This endpoint will be publicly accessible." - ) - try: - google.deploy( - model_ref, - model_method=method, - wandb_project=project, - auth_entity=auth_entity, - gcp_project=gcp_project, - service_account=service_account, - ) - except ValueError as e: - if os.getenv("DEBUG") == "true": - raise e - else: - raise click.ClickException( - str(e) + "\nRun with DEBUG=true to see full exception." - ) - print("Model deployed") - - -@deploy.command(help="Deploy to Modal Labs.") -@click.argument("model_ref") -@click.option("--project", help="W&B project name.") -@click.option( - "--auth-entity", help="Enforce authorization for this entity using W&B API keys." -) -@click.option("--dev", is_flag=True, help="Run the function locally.") -def modal(model_ref: str, project: str, auth_entity: str, dev: bool = False) -> None: - from weave.deploy import modal as mdp - - if dev: - print(f"Developing model {model_ref}...") - mdp.develop(model_ref, auth_entity=auth_entity) - return - print(f"Deploying model {model_ref}...") - if auth_entity is None: - print( - "WARNING: No --auth-entity specified. This endpoint will be publicly accessible." - ) - try: - mdp.deploy(model_ref, wandb_project=project, auth_entity=auth_entity) - except ValueError as e: - if os.getenv("DEBUG") == "true": - raise e - else: - raise click.ClickException( - str(e) + "\nRun with DEBUG=true to see full exception." - ) - print("Model deployed") - - -if __name__ == "__main__": - cli() diff --git a/weave/trace/serve_fastapi.py b/weave/trace/serve_fastapi.py deleted file mode 100644 index f40dc58e755e..000000000000 --- a/weave/trace/serve_fastapi.py +++ /dev/null @@ -1,126 +0,0 @@ -import datetime -import inspect -import typing -from typing import Annotated, Optional - -from fastapi import Depends, FastAPI, Header, HTTPException -from fastapi.security import HTTPBasic, HTTPBasicCredentials -from weave_query import cache, op_args, pyfunc_type_util, weave_pydantic # type: ignore - -from weave.trace import errors -from weave.trace.op import Op, is_op -from weave.trace.refs import ObjectRef -from weave.wandb_interface.wandb_api import WandbApiAsync - -key_cache: cache.LruTimeWindowCache[str, typing.Optional[bool]] = ( - cache.LruTimeWindowCache(datetime.timedelta(minutes=5)) -) - -api: Optional[WandbApiAsync] = None - - -def wandb_auth( - entity: str, -) -> typing.Callable[ - [typing.Optional[str]], typing.Coroutine[typing.Any, typing.Any, bool] -]: - async def auth_inner(key: Annotated[Optional[str], Depends(api_key)]) -> bool: - global api - if api is None: - api = WandbApiAsync() - if key is None: - raise HTTPException(status_code=401, detail="Missing API Key") - if len(key.split("-")[-1]) != 40: - raise HTTPException(status_code=401, detail="Invalid API Key") - - authed = key_cache.get(key) - if isinstance(authed, bool): - return authed - authed = await api.can_access_entity(entity, api_key=key) - if not authed: - raise HTTPException(status_code=403, detail="Permission Denied") - key_cache.set(key, authed) - return authed - - return auth_inner - - -def api_key( - credentials: Annotated[ - Optional[HTTPBasicCredentials], - Depends( - HTTPBasic( - auto_error=False, - description="Set your username to api and password to a W&B API Key", - ) - ), - ], - x_wandb_api_key: Annotated[ - Optional[str], Header(description="Optional W&B API Key") - ] = None, -) -> Optional[str]: - if x_wandb_api_key: - return x_wandb_api_key - elif credentials and credentials.password: - return credentials.password - else: - return None - - -def object_method_app( - obj_ref: ObjectRef, - method_name: typing.Optional[str] = None, - auth_entity: typing.Optional[str] = None, -) -> FastAPI: - obj = obj_ref.get() - - attrs: dict[str, Op] = {attr: getattr(obj, attr) for attr in dir(obj)} - op_attrs = {k: v for k, v in attrs.items() if is_op(v)} - - if not op_attrs: - raise ValueError("No ops found on object") - - if method_name is None: - if len(op_attrs) > 1: - raise ValueError( - "Multiple ops found on object ({}), must specify method_name argument".format( - ", ".join(op_attrs) - ) - ) - method_name = next(iter(op_attrs)) - - if (method := getattr(obj, method_name, None)) is None: - raise ValueError(f"Method {method_name} not found") - - if not is_op(unbound_method := method.__func__): - raise ValueError(f"Expected an op, got {unbound_method}") - - try: - args = pyfunc_type_util.determine_input_type(unbound_method) - except errors.WeaveDefinitionError as e: - raise ValueError( - f"Type for model's method '{method_name}' could not be determined. Did you annotate it with Python types? {e}" - ) - if not isinstance(args, op_args.OpNamedArgs): - raise TypeError("predict op must have named args") - - arg_types = args.weave_type().property_types - del arg_types["self"] - - Item = weave_pydantic.weave_type_to_pydantic(arg_types, name="Item") - - dependencies = [] - if auth_entity: - dependencies.append(Depends(wandb_auth(auth_entity))) - - app = FastAPI(dependencies=dependencies) - - @app.post(f"/{method_name}", summary=method_name) - async def method_route(item: Item) -> dict: # type: ignore - if inspect.iscoroutinefunction(method): - result = await method(**item.dict()) # type: ignore - else: - result = method(**item.dict()) # type: ignore - return {"result": result} - - return app