fix: add capi to pip and make it fallback (#1904)

Author: alec-flowers

deploy/sdk/src/dynamo/sdk/__init__.py

@@ -23,6 +23,7 @@
 from dynamo.sdk.core.lib import DYNAMO_IMAGE, depends, liveness, readiness, service
 from dynamo.sdk.core.protocol.interface import AbstractService
 from dynamo.sdk.lib.decorators import async_on_start, on_shutdown
+from dynamo.sdk.lib.utils import get_capi_library_path
 
 dynamo_context: dict[str, Any] = {}
 
@@ -39,4 +40,5 @@
     "abstract_endpoint",
     "liveness",
     "readiness",
+    "get_capi_library_path",
 ]

deploy/sdk/src/dynamo/sdk/lib/utils.py

@@ -132,3 +132,28 @@ def upload_graph(
     )
     if resp.status_code not in (200, 201, 204):
         raise RuntimeError(f"Failed to upload graph artifact: {resp.text}")
+
+
+def get_capi_library_path() -> str:
+    """
+    Get the path to the libdynamo_llm_capi.so library.
+
+    First checks the VLLM_KV_CAPI_PATH environment variable.
+    If not set, returns the path where the library is installed by the wheel.
+
+    Returns:
+        The path to the library.
+    """
+    # First check environment variable
+    env_path = os.environ.get("VLLM_KV_CAPI_PATH")
+    if env_path:
+        return env_path
+
+    # Fall back to the installed location
+    # The library is installed at dynamo/sdk/cli/bin/libdynamo_llm_capi.so
+    import dynamo.sdk
+
+    sdk_path = os.path.dirname(dynamo.sdk.__file__)
+    lib_path = os.path.join(sdk_path, "cli", "bin", "libdynamo_llm_capi.so")
+
+    return lib_path

hatch_build.py

@@ -28,4 +28,5 @@ def initialize(self, version, build_data):
                 f"{bin_path}/http": "dynamo/sdk/cli/bin/http",
                 f"{bin_path}/metrics": "dynamo/sdk/cli/bin/metrics",
                 f"{bin_path}/mock_worker": "dynamo/sdk/cli/bin/mock_worker",
+                f"{bin_path}/libdynamo_llm_capi.so": "dynamo/sdk/cli/bin/libdynamo_llm_capi.so",
             }

launch/dynamo-run/src/subprocess/vllm_inc.py

@@ -36,6 +36,7 @@
 )
 from dynamo.runtime import DistributedRuntime, dynamo_worker
 from dynamo.runtime.logging import configure_dynamo_logging
+from dynamo.sdk.lib.utils import get_capi_library_path
 
 # Only used if you run it manually from the command line
 DEFAULT_ENDPOINT = "dyn://dynamo.backend.generate"
@@ -208,7 +209,7 @@ async def init(runtime: DistributedRuntime, config: Config):
 
     _check_and_set_env_value("VLLM_WORKER_ID", str(endpoint.lease_id()))
     _check_and_set_env_value(
-        "VLLM_KV_CAPI_PATH", "libdynamo_llm_capi.so", allow_override=True
+        "VLLM_KV_CAPI_PATH", get_capi_library_path(), allow_override=True
     )
     _check_and_set_env_value("VLLM_KV_NAMESPACE", config.namespace)
     _check_and_set_env_value("VLLM_KV_COMPONENT", config.component)