[Serve] Type hints for API (#10205)

python/ray/serve/api.py

@@ -8,11 +8,13 @@ from ray.serve.handle import RayServeHandle
 from ray.serve.utils import (block_until_http_ready, format_actor_name)
 from ray.serve.exceptions import RayServeException
 from ray.serve.config import BackendConfig, ReplicaConfig
+from ray.actor import ActorHandle
+from typing import Any, Callable, Dict, List, Optional, Type, Union
 
 controller = None
 
 
-def _get_controller():
+def _get_controller() -> ActorHandle:
     """Used for internal purpose because using just import serve.global_state
     will always reference the original None object.
     """
@@ -23,7 +25,7 @@ def _get_controller():
     return controller
 
 
-def _ensure_connected(f):
+def _ensure_connected(f: Callable) -> Callable:
     @wraps(f)
     def check(*args, **kwargs):
         _get_controller()
@@ -32,7 +34,7 @@ def _ensure_connected(f):
     return check
 
 
-def accept_batch(f):
+def accept_batch(f: Callable) -> Callable:
     """Annotation to mark a serving function that batch is accepted.
 
     This annotation need to be used to mark a function expect all arguments
@@ -54,10 +56,10 @@ def accept_batch(f):
     return f
 
 
-def init(name=None,
-         http_host=DEFAULT_HTTP_HOST,
-         http_port=DEFAULT_HTTP_PORT,
-         _http_middlewares=[]):
+def init(name: Optional[str] = None,
+         http_host: str = DEFAULT_HTTP_HOST,
+         http_port: int = DEFAULT_HTTP_PORT,
+         _http_middlewares: List[Any] = []) -> None:
     """Initialize or connect to a serve cluster.
 
     If serve cluster is already initialized, this function will just return.
@@ -109,7 +111,7 @@ def init(name=None,
 
 
 @_ensure_connected
-def shutdown():
+def shutdown() -> None:
     """Completely shut down the connected Serve instance.
 
     Shuts down all processes and deletes all state associated with the Serve
@@ -122,11 +124,11 @@ def shutdown():
 
 
 @_ensure_connected
-def create_endpoint(endpoint_name,
+def create_endpoint(endpoint_name: str,
                     *,
-                    backend=None,
-                    route=None,
-                    methods=["GET"]):
+                    backend: str = None,
+                    route: Optional[str] = None,
+                    methods: List[str] = ["GET"]) -> None:
     """Create a service endpoint given route_expression.
 
     Args:
@@ -179,7 +181,7 @@ def create_endpoint(endpoint_name,
 
 
 @_ensure_connected
-def delete_endpoint(endpoint):
+def delete_endpoint(endpoint: str) -> None:
     """Delete the given endpoint.
 
     Does not delete any associated backends.
@@ -188,7 +190,7 @@ def delete_endpoint(endpoint):
 
 
 @_ensure_connected
-def list_endpoints():
+def list_endpoints() -> Dict[str, Dict[str, Any]]:
     """Returns a dictionary of all registered endpoints.
 
     The dictionary keys are endpoint names and values are dictionaries
@@ -198,7 +200,8 @@ def list_endpoints():
 
 
 @_ensure_connected
-def update_backend_config(backend_tag, config_options):
+def update_backend_config(backend_tag: str,
+                          config_options: Dict[str, Any]) -> None:
     """Update a backend configuration for a backend tag.
 
     Keys not specified in the passed will be left unchanged.
@@ -225,7 +228,7 @@ def update_backend_config(backend_tag, config_options):
 
 
 @_ensure_connected
-def get_backend_config(backend_tag):
+def get_backend_config(backend_tag: str):
     """Get the backend configuration for a backend tag.
 
     Args:
@@ -235,11 +238,11 @@ def get_backend_config(backend_tag):
 
 
 @_ensure_connected
-def create_backend(backend_tag,
-                   func_or_class,
-                   *actor_init_args,
-                   ray_actor_options=None,
-                   config=None):
+def create_backend(backend_tag: str,
+                   func_or_class: Union[Callable, Type[Callable]],
+                   *actor_init_args: Any,
+                   ray_actor_options: Optional[Dict] = None,
+                   config: Optional[Dict[str, Any]] = None) -> None:
     """Create a backend with the provided tag.
 
     The backend will serve requests with func_or_class.
@@ -286,7 +289,7 @@ def create_backend(backend_tag,
 
 
 @_ensure_connected
-def list_backends():
+def list_backends() -> Dict[str, Dict[str, Any]]:
     """Returns a dictionary of all registered backends.
 
     Dictionary maps backend tags to backend configs.
@@ -295,7 +298,7 @@ def list_backends():
 
 
 @_ensure_connected
-def delete_backend(backend_tag):
+def delete_backend(backend_tag: str) -> None:
     """Delete the given backend.
 
     The backend must not currently be used by any endpoints.
@@ -304,7 +307,8 @@ def delete_backend(backend_tag):
 
 
 @_ensure_connected
-def set_traffic(endpoint_name, traffic_policy_dictionary):
+def set_traffic(endpoint_name: str,
+                traffic_policy_dictionary: Dict[str, float]) -> None:
     """Associate a service endpoint with traffic policy.
 
     Example:
@@ -325,7 +329,8 @@ def set_traffic(endpoint_name, traffic_policy_dictionary):
 
 
 @_ensure_connected
-def shadow_traffic(endpoint_name, backend_tag, proportion):
+def shadow_traffic(endpoint_name: str, backend_tag: str,
+                   proportion: float) -> None:
     """Shadow traffic from an endpoint to a backend.
 
     The specified proportion of requests will be duplicated and sent to the
@@ -350,7 +355,7 @@ def shadow_traffic(endpoint_name, backend_tag, proportion):
 
 
 @_ensure_connected
-def get_handle(endpoint_name, missing_ok=False):
+def get_handle(endpoint_name: str, missing_ok: bool = False) -> RayServeHandle:
     """Retrieve RayServeHandle for service endpoint to invoke it from Python.
 
     Args: