diff --git a/.dockerignore b/.dockerignore index 63811e87aa0f..1eeea4f18f41 100644 --- a/.dockerignore +++ b/.dockerignore @@ -19,7 +19,7 @@ **/*.plan **/.cache/* **/*onnx* -# Engine must be allowed because code contains triton_distributed_engine.py +# Engine must be allowed because code contains dynemo_engine.py **/*tensorrtllm_engines* **/*tensorrtllm_models* **/*tensorrtllm_checkpoints* diff --git a/.github/workflows/pre-merge.yml b/.github/workflows/pre-merge.yml index ab311c4c6a55..c67f6d569300 100644 --- a/.github/workflows/pre-merge.yml +++ b/.github/workflows/pre-merge.yml @@ -22,25 +22,6 @@ on: jobs: - # icp_validation: - # runs-on: ubuntu-latest - # container: - # image: ghcr.io/triton-inference-server/triton3/python_ci:0.1.9 - # env: - # BUILD_NUMBER: ${{ github.job }} - # CUDA_VISIBLE_DEVICES: -1 - # PATH: /opt/tritonserver/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/ucx/bin:/bin:/sbin:/usr/bin:/usr/sbin:/usr/local/bin:/usr/local/mpi/bin:/usr/local/sbin - # volumes: - # - ${{ github.workspace }}:/workspace - # permissions: - # contents: read - # packages: read - # steps: - # - uses: actions/checkout@v4 - # - run: ./icp/protos/gen_python.sh - # - run: pytest --verbose icp - # timeout-minutes: 3 - pre-commit: runs-on: ubuntu-latest permissions: @@ -52,41 +33,3 @@ jobs: timeout-minutes: 3 - - # providers_validation: - # runs-on: ubuntu-latest - # container: - # image: ghcr.io/triton-inference-server/triton3/python_ci:0.1.9 - # env: - # BUILD_NUMBER: ${{ github.job }} - # CUDA_VISIBLE_DEVICES: -1 - # PATH: /opt/tritonserver/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/ucx/bin:/bin:/sbin:/usr/bin:/usr/sbin:/usr/local/bin:/usr/local/mpi/bin:/usr/local/sbin - # PROTO_OUT: /python/icp/protos - # volumes: - # - ${{ github.workspace }}:/workspace - # permissions: - # contents: read - # packages: read - # steps: - # - uses: actions/checkout@v4 - # - run: pytest --verbose providers - - # worker_validation: - # runs-on: ubuntu-latest - # container: - # image: ghcr.io/triton-inference-server/triton3/python_ci:0.1.9 - # env: - # BUILD_NUMBER: ${{ github.job }} - # CUDA_VISIBLE_DEVICES: -1 - # PATH: /opt/tritonserver/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/ucx/bin:/bin:/sbin:/usr/bin:/usr/sbin:/usr/local/bin:/usr/local/mpi/bin:/usr/local/sbin - # PROTO_OUT: /python/icp/protos - # volumes: - # - ${{ github.workspace }}:/workspace - # permissions: - # contents: read - # packages: read - # steps: - # - uses: actions/checkout@v4 - # - run: ./icp/protos/gen_python.sh - # - run: pytest -p no:warnings --verbose worker/python/tests - # timeout-minutes: 2 diff --git a/ATTRIBUTIONS.md b/ATTRIBUTIONS.md index b4bcee75fff9..08d13878d87b 100644 --- a/ATTRIBUTIONS.md +++ b/ATTRIBUTIONS.md @@ -17,7 +17,7 @@ limitations under the License. # Open Source License Attribution - Triton Distributed uses Open Source components. You can find the details of these open-source projects along with license information below. + Dynemo uses Open Source components. You can find the details of these open-source projects along with license information below. We are grateful to the developers for their contributions to open source and acknowledge these below. ## nats-py - [Apache License 2.0](https://github.com/nats-io/nats.py/blob/main/LICENSE) diff --git a/README.md b/README.md index 1d8fb43bb8e4..0c268e271e72 100644 --- a/README.md +++ b/README.md @@ -71,7 +71,7 @@ The run script offers a few common workflows: 1. Running a command in a container and exiting. ``` -./container/run.sh -- python3 -c "import triton_distributed.runtime; help(triton_distributed.runtime)" +./container/run.sh -- python3 -c "import dynemo.runtime; help(dynemo.runtime)" ``` 2. Starting an interactive shell. diff --git a/applications/llm/count/Cargo.lock b/applications/llm/count/Cargo.lock index 0204f8db60ab..ca088fa16b8f 100644 --- a/applications/llm/count/Cargo.lock +++ b/applications/llm/count/Cargo.lock @@ -737,6 +737,8 @@ version = "0.1.0" dependencies = [ "axum 0.6.20", "clap", + "dynemo-llm", + "dynemo-runtime", "opentelemetry", "opentelemetry-prometheus", "prometheus", @@ -747,8 +749,6 @@ dependencies = [ "thiserror 1.0.69", "tokio", "tracing", - "triton-distributed-llm", - "triton-distributed-runtime", ] [[package]] @@ -1024,6 +1024,99 @@ dependencies = [ "syn 2.0.98", ] +[[package]] +name = "dynemo-llm" +version = "0.2.1" +dependencies = [ + "anyhow", + "async-openai", + "async-stream", + "async-trait", + "axum 0.8.1", + "bindgen", + "blake3", + "bs62", + "bytes", + "chrono", + "cmake", + "derive_builder", + "dynemo-runtime", + "either", + "erased-serde", + "futures", + "galil-seiferas", + "indexmap 2.7.1", + "itertools 0.14.0", + "libc", + "minijinja", + "minijinja-contrib", + "prometheus", + "pyo3", + "regex", + "semver", + "serde", + "serde-pickle", + "serde_json", + "serde_repr", + "strum", + "thiserror 2.0.11", + "tokenizers", + "tokio", + "tokio-stream", + "tokio-util", + "toktrie", + "toktrie_hf_tokenizers", + "tracing", + "unicode-segmentation", + "uuid", + "validator", + "xxhash-rust", +] + +[[package]] +name = "dynemo-runtime" +version = "0.2.1" +dependencies = [ + "anyhow", + "async-nats", + "async-once-cell", + "async-stream", + "async-trait", + "async_zmq", + "blake3", + "bytes", + "chrono", + "derive-getters", + "derive_builder", + "educe", + "either", + "etcd-client", + "figment", + "futures", + "humantime", + "local-ip-address", + "log", + "nid", + "nix", + "nuid", + "once_cell", + "prometheus", + "rand", + "regex", + "serde", + "serde_json", + "socket2", + "thiserror 1.0.69", + "tokio", + "tokio-stream", + "tokio-util", + "tracing", + "tracing-subscriber", + "uuid", + "validator", + "xxhash-rust", +] + [[package]] name = "ed25519" version = "2.2.3" @@ -4232,99 +4325,6 @@ dependencies = [ "tracing-serde", ] -[[package]] -name = "triton-distributed-llm" -version = "0.2.1" -dependencies = [ - "anyhow", - "async-openai", - "async-stream", - "async-trait", - "axum 0.8.1", - "bindgen", - "blake3", - "bs62", - "bytes", - "chrono", - "cmake", - "derive_builder", - "either", - "erased-serde", - "futures", - "galil-seiferas", - "indexmap 2.7.1", - "itertools 0.14.0", - "libc", - "minijinja", - "minijinja-contrib", - "prometheus", - "pyo3", - "regex", - "semver", - "serde", - "serde-pickle", - "serde_json", - "serde_repr", - "strum", - "thiserror 2.0.11", - "tokenizers", - "tokio", - "tokio-stream", - "tokio-util", - "toktrie", - "toktrie_hf_tokenizers", - "tracing", - "triton-distributed-runtime", - "unicode-segmentation", - "uuid", - "validator", - "xxhash-rust", -] - -[[package]] -name = "triton-distributed-runtime" -version = "0.2.1" -dependencies = [ - "anyhow", - "async-nats", - "async-once-cell", - "async-stream", - "async-trait", - "async_zmq", - "blake3", - "bytes", - "chrono", - "derive-getters", - "derive_builder", - "educe", - "either", - "etcd-client", - "figment", - "futures", - "humantime", - "local-ip-address", - "log", - "nid", - "nix", - "nuid", - "once_cell", - "prometheus", - "rand", - "regex", - "serde", - "serde_json", - "socket2", - "thiserror 1.0.69", - "tokio", - "tokio-stream", - "tokio-util", - "tracing", - "tracing-subscriber", - "uuid", - "validator", - "xxhash-rust", -] - [[package]] name = "try-lock" version = "0.2.5" diff --git a/applications/llm/count/Cargo.toml b/applications/llm/count/Cargo.toml index 46e05cdaf61d..d05022c78e8b 100644 --- a/applications/llm/count/Cargo.toml +++ b/applications/llm/count/Cargo.toml @@ -21,8 +21,8 @@ license = "Apache-2.0" [dependencies] # local -triton-distributed-runtime = { path = "../../../lib/runtime" } -triton-distributed-llm = { path = "../../../lib/llm" } +dynemo-runtime = { path = "../../../lib/runtime" } +dynemo-llm = { path = "../../../lib/llm" } # workspace - todo @@ -40,4 +40,4 @@ rand = "0.8" axum = "0.6" [dev-dependencies] -reqwest = { version = "0.11", features = ["blocking"] } \ No newline at end of file +reqwest = { version = "0.11", features = ["blocking"] } diff --git a/applications/llm/count/README.md b/applications/llm/count/README.md index 1f7d80b5033d..43d7eae37dfb 100644 --- a/applications/llm/count/README.md +++ b/applications/llm/count/README.md @@ -8,17 +8,17 @@ the services associated with that endpoint, do some postprocessing on them, and then publish an event with the postprocessed data. ```bash -# For more details, try TRD_LOG=debug -TRD_LOG=info cargo run --bin count -- --namespace triton-init --component backend --endpoint generate +# For more details, try DYN_LOG=debug +DYN_LOG=info cargo run --bin count -- --namespace dynemo --component backend --endpoint generate -# 2025-02-26T18:45:05.467026Z INFO count: Creating unique instance of Count at triton-init/components/count/instance -# 2025-02-26T18:45:05.472146Z INFO count: Scraping service triton_init_backend_720278f8 and filtering on subject triton_init_backend_720278f8.generate +# 2025-02-26T18:45:05.467026Z INFO count: Creating unique instance of Count at dynemo/components/count/instance +# 2025-02-26T18:45:05.472146Z INFO count: Scraping service dynemo_init_backend_720278f8 and filtering on subject dynemo_init_backend_720278f8.generate # ... ``` With no matching endpoints running, you should see warnings in the logs: ```bash -2025-02-26T18:45:06.474161Z WARN count: No endpoints found matching subject triton_init_backend_720278f8.generate +2025-02-26T18:45:06.474161Z WARN count: No endpoints found matching subject dynemo_init_backend_720278f8.generate ``` To see metrics published to a matching endpoint, you can use the @@ -35,7 +35,7 @@ since the endpoint will automatically get discovered. When stats are found from the target endpoints being listened on, count will aggregate and publish some metrics as both an event and to a prometheus web server: ``` -2025-02-28T04:05:58.077901Z INFO count: Aggregated metrics: ProcessedEndpoints { endpoints: [Endpoint { name: "worker-7587884888253033398", subject: "triton_init_backend_720278f8.generate-694d951a80e06bb6", data: ForwardPassMetrics { request_active_slots: 58, request_total_slots: 100, kv_active_blocks: 77, kv_total_blocks: 100 } }, Endpoint { name: "worker-7587884888253033401", subject: "triton_init_backend_720278f8.generate-694d951a80e06bb9", data: ForwardPassMetrics { request_active_slots: 71, request_total_slots: 100, kv_active_blocks: 29, kv_total_blocks: 100 } }], worker_ids: [7587884888253033398, 7587884888253033401], load_avg: 53.0, load_std: 24.0 } +2025-02-28T04:05:58.077901Z INFO count: Aggregated metrics: ProcessedEndpoints { endpoints: [Endpoint { name: "worker-7587884888253033398", subject: "dynemo_init_backend_720278f8.generate-694d951a80e06bb6", data: ForwardPassMetrics { request_active_slots: 58, request_total_slots: 100, kv_active_blocks: 77, kv_total_blocks: 100 } }, Endpoint { name: "worker-7587884888253033401", subject: "dynemo_init_backend_720278f8.generate-694d951a80e06bb9", data: ForwardPassMetrics { request_active_slots: 71, request_total_slots: 100, kv_active_blocks: 29, kv_total_blocks: 100 } }], worker_ids: [7587884888253033398, 7587884888253033401], load_avg: 53.0, load_std: 24.0 } ``` To see the metrics being published in prometheus format, you can run: diff --git a/applications/llm/count/src/bin/mock_worker.rs b/applications/llm/count/src/bin/mock_worker.rs index 08ed2b2f704a..b0297c0cd81e 100644 --- a/applications/llm/count/src/bin/mock_worker.rs +++ b/applications/llm/count/src/bin/mock_worker.rs @@ -13,10 +13,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use rand::Rng; -use std::sync::Arc; -use triton_distributed_llm::kv_router::protocols::ForwardPassMetrics; -use triton_distributed_runtime::{ +use dynemo_llm::kv_router::protocols::ForwardPassMetrics; +use dynemo_runtime::{ logging, pipeline::{ async_trait, network::Ingress, AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, @@ -25,6 +23,8 @@ use triton_distributed_runtime::{ protocols::annotated::Annotated, stream, DistributedRuntime, Result, Runtime, Worker, }; +use rand::Rng; +use std::sync::Arc; fn main() -> Result<()> { logging::init(); @@ -69,7 +69,7 @@ async fn backend(runtime: DistributedRuntime) -> Result<()> { // we must first create a service, then we can attach one more more endpoints runtime - .namespace("triton-init")? + .namespace("dynemo")? .component("backend")? .service_builder() .create() diff --git a/applications/llm/count/src/lib.rs b/applications/llm/count/src/lib.rs index 0612bd81bed2..f439af89452f 100644 --- a/applications/llm/count/src/lib.rs +++ b/applications/llm/count/src/lib.rs @@ -20,13 +20,11 @@ use prometheus::register_gauge_vec; use serde::{Deserialize, Serialize}; use std::net::SocketAddr; -use triton_distributed_llm::kv_router::protocols::ForwardPassMetrics; -use triton_distributed_llm::kv_router::scheduler::Endpoint; -use triton_distributed_llm::kv_router::scoring::ProcessedEndpoints; +use dynemo_llm::kv_router::protocols::ForwardPassMetrics; +use dynemo_llm::kv_router::scheduler::Endpoint; +use dynemo_llm::kv_router::scoring::ProcessedEndpoints; -use triton_distributed_runtime::{ - distributed::Component, service::EndpointInfo, utils::Duration, Result, -}; +use dynemo_runtime::{distributed::Component, service::EndpointInfo, utils::Duration, Result}; /// Configuration for LLM worker load capacity metrics #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/applications/llm/count/src/main.rs b/applications/llm/count/src/main.rs index 0f7634d518eb..513a2b0b0b99 100644 --- a/applications/llm/count/src/main.rs +++ b/applications/llm/count/src/main.rs @@ -24,7 +24,7 @@ //! - KV Cache Blocks: [Active, Total] use clap::Parser; -use triton_distributed_runtime::{ +use dynemo_runtime::{ error, logging, traits::events::EventPublisher, utils::{Duration, Instant}, @@ -50,7 +50,7 @@ struct Args { endpoint: String, /// Namespace to operate in - #[arg(long, env = "TRD_NAMESPACE", default_value = "triton-init")] + #[arg(long, env = "DYN_NAMESPACE", default_value = "dynemo")] namespace: String, /// Polling interval in seconds (minimum 1 second) @@ -155,7 +155,7 @@ mod tests { #[test] fn test_namespace_from_env() { - env::set_var("TRD_NAMESPACE", "test-namespace"); + env::set_var("DYN_NAMESPACE", "test-namespace"); let args = Args::parse_from(["count", "--component", "comp", "--endpoint", "end"]); assert_eq!(args.namespace, "test-namespace"); } diff --git a/container/Dockerfile b/container/Dockerfile index a2a9e8274c6b..f0c1f2b068af 100644 --- a/container/Dockerfile +++ b/container/Dockerfile @@ -16,7 +16,7 @@ ARG BASE_IMAGE="nvcr.io/nvidia/tritonserver" ARG BASE_IMAGE_TAG="25.01-py3" -FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG} AS triton-distributed +FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG} AS dynemo # TODO: non root user by default @@ -34,7 +34,7 @@ RUN rustup toolchain install 1.85.0-x86_64-unknown-linux-gnu # Install OpenAI-compatible frontend and its dependencies from triton server # repository. These are used to have a consistent interface, schema, and FastAPI -# app between Triton Core and Triton Distributed implementations. +# app between Triton Core and Dynemo implementations. ARG OPENAI_SERVER_TAG="r25.01" RUN mkdir -p /opt/tritonserver/python && \ cd /opt/tritonserver/python && \ @@ -78,7 +78,7 @@ ARG TENSORRTLLM_SKIP_CLONE= ENV FRAMEWORK=${FRAMEWORK} RUN --mount=type=bind,source=./container/deps/requirements.tensorrtllm.txt,target=/tmp/requirements.txt \ --mount=type=bind,source=./container/deps/clone_tensorrtllm.sh,target=/tmp/clone_tensorrtllm.sh \ - if [[ "$FRAMEWORK" == "TENSORRTLLM" ]] ; then pip install --timeout=2000 -r /tmp/requirements.txt; if [ ${TENSORRTLLM_SKIP_CLONE} -ne 1 ] ; then /tmp/clone_tensorrtllm.sh --tensorrtllm-backend-repo-tag ${TENSORRTLLM_BACKEND_REPO_TAG} --tensorrtllm-backend-rebuild ${TENSORRTLLM_BACKEND_REBUILD} --triton-llm-path /opt/triton/llm_binding ; fi ; fi + if [[ "$FRAMEWORK" == "TENSORRTLLM" ]] ; then pip install --timeout=2000 -r /tmp/requirements.txt; if [ ${TENSORRTLLM_SKIP_CLONE} -ne 1 ] ; then /tmp/clone_tensorrtllm.sh --tensorrtllm-backend-repo-tag ${TENSORRTLLM_BACKEND_REPO_TAG} --tensorrtllm-backend-rebuild ${TENSORRTLLM_BACKEND_REBUILD} --triton-llm-path /opt/dynemo/llm_binding ; fi ; fi RUN --mount=type=bind,source=./container/deps/requirements.standard.txt,target=/tmp/requirements.txt \ @@ -106,7 +106,7 @@ ENV VLLM_GENERATE_WORKERS=${VLLM_FRAMEWORK:+1} ENV VLLM_BASELINE_TP_SIZE=${VLLM_FRAMEWORK:+1} ENV VLLM_CONTEXT_TP_SIZE=${VLLM_FRAMEWORK:+1} ENV VLLM_GENERATE_TP_SIZE=${VLLM_FRAMEWORK:+1} -ENV VLLM_KV_CAPI_PATH="/opt/triton/llm_binding/lib/libtriton_llm_capi.so" +ENV VLLM_KV_CAPI_PATH="/opt/dynemo/llm_binding/lib/libdynemo_llm_capi.so" ENV PYTHONUNBUFFERED=1 # Install NATS - pointing toward NATS github instead of binaries.nats.dev due to server instability @@ -159,27 +159,27 @@ COPY lib/bindings /workspace/lib/bindings RUN cd lib/bindings/c/ && \ cargo build --release --locked && cargo doc --no-deps -# Install uv, create virtualenv for general use, and build triton_distributed wheel +# Install uv, create virtualenv for general use, and build dynemo wheel COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/ -RUN mkdir /opt/triton && \ - uv venv /opt/triton/venv --python 3.12 && \ - source /opt/triton/venv/bin/activate && \ +RUN mkdir /opt/dynemo && \ + uv venv /opt/dynemo/venv --python 3.12 && \ + source /opt/dynemo/venv/bin/activate && \ uv build --wheel --out-dir /workspace/dist && \ - uv pip install /workspace/dist/triton_distributed*cp312*.whl + uv pip install /workspace/dist/dynemo*cp312*.whl # Package the bindings -RUN mkdir -p /opt/triton/bindings/wheels && \ - mkdir /opt/triton/bindings/lib && \ - cp dist/triton_distributed*cp312*.whl /opt/triton/bindings/wheels/. && \ - cp lib/bindings/c/target/release/libtriton_distributed_llm_capi.so /opt/triton/bindings/lib/. && \ - cp -r lib/bindings/c/include /opt/triton/bindings/. +RUN mkdir -p /opt/dynemo/bindings/wheels && \ + mkdir /opt/dynemo/bindings/lib && \ + cp dist/dynemo*cp312*.whl /opt/dynemo/bindings/wheels/. && \ + cp lib/bindings/c/target/release/libdynemo_llm_capi.so /opt/dynemo/bindings/lib/. && \ + cp -r lib/bindings/c/include /opt/dynemo/bindings/. -# Install triton_distributed_runtime and triton_distributed_llm wheels globally in container for tests that +# Install dynemo.runtime and dynemo.llm wheels globally in container for tests that # currently run without virtual environment activated. # TODO: In future, we may use a virtualenv for everything and remove this. -RUN pip install /opt/triton/bindings/wheels/triton_distributed*cp312*.whl +RUN pip install /opt/dynemo/bindings/wheels/dynemo*cp312*.whl -# Copy everything in after install steps to avoid re-running build/install +# Copy everything in after ginstall steps to avoid re-running build/install # commands on unrelated changes in other dirs. COPY . /workspace diff --git a/container/Dockerfile.vllm b/container/Dockerfile.vllm index 1ed100c56ab1..9e8e085b4837 100644 --- a/container/Dockerfile.vllm +++ b/container/Dockerfile.vllm @@ -24,17 +24,17 @@ ENV PATH=/usr/local/bin/etcd/:$PATH # Install uv and create virtualenv COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/ -RUN mkdir /opt/triton && \ - uv venv /opt/triton/venv --python 3.12 +RUN mkdir /opt/dynemo && \ + uv venv /opt/dynemo/venv --python 3.12 # Activate virtual environment -ENV VIRTUAL_ENV=/opt/triton/venv +ENV VIRTUAL_ENV=/opt/dynemo/venv ENV PATH="${VIRTUAL_ENV}/bin:${PATH}" # Install patched vllm - keep this early in Dockerfile to avoid # rebuilds from unrelated source code changes ARG VLLM_REF="v0.7.2" -ARG VLLM_PATCH="vllm_${VLLM_REF}-triton-kv-disagg-patch.patch" +ARG VLLM_PATCH="vllm_${VLLM_REF}-dynemo-kv-disagg-patch.patch" RUN --mount=type=bind,source=./container/deps/,target=/tmp/deps \ bash /tmp/deps/vllm/install.sh --patch /tmp/deps/vllm/${VLLM_PATCH} --ref ${VLLM_REF} --install-cmd "uv pip install --editable" --use-precompiled --installation-dir /opt/vllm @@ -100,25 +100,25 @@ COPY lib/bindings /workspace/lib/bindings RUN cd lib/bindings/c && \ cargo build --release --locked && cargo doc --no-deps -# Build triton_distributed wheel -RUN source /opt/triton/venv/bin/activate && \ +# Build dynemo wheel +RUN source /opt/dynemo/venv/bin/activate && \ uv build --wheel --out-dir /workspace/dist && \ - uv pip install /workspace/dist/triton_distributed*cp312*.whl + uv pip install /workspace/dist/dynemo*cp312*.whl # Package the bindings -RUN mkdir -p /opt/triton/bindings/wheels && \ - mkdir /opt/triton/bindings/lib && \ - cp dist/triton_distributed*cp312*.whl /opt/triton/bindings/wheels/. && \ - cp lib/bindings/c/target/release/libtriton_distributed_llm_capi.so /opt/triton/bindings/lib/. && \ - cp -r lib/bindings/c/include /opt/triton/bindings/. +RUN mkdir -p /opt/dynemo/bindings/wheels && \ + mkdir /opt/dynemo/bindings/lib && \ + cp dist/dynemo*cp312*.whl /opt/dynemo/bindings/wheels/. && \ + cp lib/bindings/c/target/release/libdynemo_llm_capi.so /opt/dynemo/bindings/lib/. && \ + cp -r lib/bindings/c/include /opt/dynemo/bindings/. -# Tell vllm to use the Triton LLM C API for KV Cache Routing -ENV VLLM_KV_CAPI_PATH="/opt/triton/bindings/lib/libtriton_distributed_llm_capi.so" +# Tell vllm to use the Dynemo LLM C API for KV Cache Routing +ENV VLLM_KV_CAPI_PATH="/opt/dynemo/bindings/lib/libdynemo_llm_capi.so" # FIXME: Copy more specific folders in for dev/debug after directory restructure COPY . /workspace -# FIXME: May want a modification with triton-distributed banner on entry +# FIXME: May want a modification with dynemo-distributed banner on entry ENTRYPOINT ["/opt/nvidia/nvidia_entrypoint.sh"] CMD [] @@ -136,10 +136,10 @@ RUN apt update -y && \ echo "set -g mouse on" >> /root/.tmux.conf # Set environment variables -ENV VIRTUAL_ENV=/opt/triton/venv +ENV VIRTUAL_ENV=/opt/dynemo/venv ENV PATH="${VIRTUAL_ENV}/bin:${PATH}" ENV RAPIDS_LIBUCX_PREFER_SYSTEM_LIBRARY=true -ENV VLLM_KV_CAPI_PATH="/opt/triton/bindings/lib/libtriton_distributed_llm_capi.so" +ENV VLLM_KV_CAPI_PATH="/opt/dynemo/bindings/lib/libdynemo_llm_capi.so" # Copy binaries COPY --from=dev /usr/local/bin/http /usr/local/bin/http @@ -166,7 +166,7 @@ COPY examples/python_rs/llm/vllm /workspace/examples/python_rs/llm/vllm WORKDIR /workspace -# FIXME: May want a modification with triton-distributed banner on entry +# FIXME: May want a modification with dynemo-distributed banner on entry ENTRYPOINT ["/opt/nvidia/nvidia_entrypoint.sh"] CMD [] diff --git a/container/Dockerfile.vllm_nixl b/container/Dockerfile.vllm_nixl index c40451b6fa76..ce4c06ba42c6 100644 --- a/container/Dockerfile.vllm_nixl +++ b/container/Dockerfile.vllm_nixl @@ -150,17 +150,17 @@ ENV PATH=/usr/local/bin/etcd/:$PATH # Install uv and create virtualenv COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/ -RUN mkdir /opt/triton && \ - uv venv /opt/triton/venv --python 3.12 +RUN mkdir /opt/dynemo && \ + uv venv /opt/dynemo/venv --python 3.12 # Activate virtual environment -ENV VIRTUAL_ENV=/opt/triton/venv +ENV VIRTUAL_ENV=/opt/dynemo/venv ENV PATH="${VIRTUAL_ENV}/bin:${PATH}" # Install patched vllm - keep this early in Dockerfile to avoid # rebuilds from unrelated source code changes ARG VLLM_REF="v0.7.2" -ARG VLLM_PATCH="vllm_${VLLM_REF}-triton-kv-disagg-patch.patch" +ARG VLLM_PATCH="vllm_${VLLM_REF}-dynemo-kv-disagg-patch.patch" RUN --mount=type=bind,source=./container/deps/,target=/tmp/deps \ bash /tmp/deps/vllm/install.sh --patch /tmp/deps/vllm/${VLLM_PATCH} --ref ${VLLM_REF} --install-cmd "uv pip install --editable" --use-precompiled --installation-dir /opt/vllm @@ -225,25 +225,25 @@ COPY lib/bindings /workspace/lib/bindings RUN cd lib/bindings/c && \ cargo build --release --locked && cargo doc --no-deps -# Build triton_distributed wheel -RUN source /opt/triton/venv/bin/activate && \ +# Build dynemo wheel +RUN source /opt/dynemo/venv/bin/activate && \ uv build --wheel --out-dir /workspace/dist && \ - uv pip install /workspace/dist/triton_distributed*cp312*.whl + uv pip install /workspace/dist/dynemo*cp312*.whl # Package the bindings -RUN mkdir -p /opt/triton/bindings/wheels && \ - mkdir /opt/triton/bindings/lib && \ - cp dist/triton_distributed*cp312*.whl /opt/triton/bindings/wheels/. && \ - cp lib/bindings/c/target/release/libtriton_distributed_llm_capi.so /opt/triton/bindings/lib/. && \ - cp -r lib/bindings/c/include /opt/triton/bindings/. +RUN mkdir -p /opt/dynemo/bindings/wheels && \ + mkdir /opt/dynemo/bindings/lib && \ + cp dist/dynemo*cp312*.whl /opt/dynemo/bindings/wheels/. && \ + cp lib/bindings/c/target/release/libdynemo_llm_capi.so /opt/dynemo/bindings/lib/. && \ + cp -r lib/bindings/c/include /opt/dynemo/bindings/. -# Tell vllm to use the Triton LLM C API for KV Cache Routing -ENV VLLM_KV_CAPI_PATH="/opt/triton/bindings/lib/libtriton_distributed_llm_capi.so" +# Tell vllm to use the Dynemo LLM C API for KV Cache Routing +ENV VLLM_KV_CAPI_PATH="/opt/dynemo/bindings/lib/libdynemo_llm_capi.so" # FIXME: Copy more specific folders in for dev/debug after directory restructure COPY . /workspace -# FIXME: May want a modification with triton-distributed banner on entry +# FIXME: May want a modification with dynemo-distributed banner on entry ENTRYPOINT ["/opt/nvidia/nvidia_entrypoint.sh"] CMD [] @@ -261,10 +261,10 @@ RUN apt update -y && \ echo "set -g mouse on" >> /root/.tmux.conf # Set environment variables -ENV VIRTUAL_ENV=/opt/triton/venv +ENV VIRTUAL_ENV=/opt/dynemo/venv ENV PATH="${VIRTUAL_ENV}/bin:${PATH}" ENV RAPIDS_LIBUCX_PREFER_SYSTEM_LIBRARY=true -ENV VLLM_KV_CAPI_PATH="/opt/triton/bindings/lib/libtriton_distributed_llm_capi.so" +ENV VLLM_KV_CAPI_PATH="/opt/dynemo/bindings/lib/libdynemo_llm_capi.so" # Copy binaries COPY --from=dev /usr/local/bin/http /usr/local/bin/http @@ -291,7 +291,7 @@ COPY examples/python_rs/llm/vllm_nixl /workspace/examples/python_rs/llm/vllm_nix WORKDIR /workspace -# FIXME: May want a modification with triton-distributed banner on entry +# FIXME: May want a modification with dynemo-distributed banner on entry ENTRYPOINT ["/opt/nvidia/nvidia_entrypoint.sh"] CMD [] diff --git a/container/deps/clone_tensorrtllm.sh b/container/deps/clone_tensorrtllm.sh index ff3a7e9c1bd5..b5aecc490203 100755 --- a/container/deps/clone_tensorrtllm.sh +++ b/container/deps/clone_tensorrtllm.sh @@ -16,7 +16,7 @@ TENSORRTLLM_BACKEND_REPO_TAG= TENSORRTLLM_BACKEND_REBUILD= -TRITON_LLM_PATH= +DYNEMO_LLM_PATH= GIT_TOKEN= GIT_REPO= @@ -43,9 +43,9 @@ get_options() { missing_requirement $1 fi ;; - --triton-llm-path) + --dynemo-llm-path) if [ "$2" ]; then - TRITON_LLM_PATH=$2 + DYNEMO_LLM_PATH=$2 shift else missing_requirement $1 @@ -147,9 +147,9 @@ if [ ! -z ${TENSORRTLLM_BACKEND_REBUILD} ]; then # Build the backend (cd inflight_batcher_llm/src \ - && cmake -DCMAKE_INSTALL_PREFIX:PATH=`pwd`/install -DUSE_CXX11_ABI=1 -DTRITON_LLM_PATH=$TRITON_LLM_PATH .. \ + && cmake -DCMAKE_INSTALL_PREFIX:PATH=`pwd`/install -DUSE_CXX11_ABI=1 -DDYNEMO_LLM_PATH=$DYNEMO_LLM_PATH .. \ && make install \ - && cp libtriton_tensorrtllm.so /opt/tritonserver/backends/tensorrtllm/ \ + && cp libdynemo_tensorrtllm.so /opt/tritonserver/backends/tensorrtllm/ \ && cp trtllmExecutorWorker /opt/tritonserver/backends/tensorrtllm/ \ ) fi diff --git a/container/deps/vllm/vllm_v0.7.2-triton-kv-disagg-patch.patch b/container/deps/vllm/vllm_v0.7.2-dynemo-kv-disagg-patch.patch similarity index 88% rename from container/deps/vllm/vllm_v0.7.2-triton-kv-disagg-patch.patch rename to container/deps/vllm/vllm_v0.7.2-dynemo-kv-disagg-patch.patch index 66e47253eb96..21f270746b8b 100644 --- a/container/deps/vllm/vllm_v0.7.2-triton-kv-disagg-patch.patch +++ b/container/deps/vllm/vllm_v0.7.2-dynemo-kv-disagg-patch.patch @@ -31,7 +31,7 @@ index 9ba49757..a2f88854 100644 f"and `kv_both`") - if self.kv_connector is not None and self.kv_role is None: -+ if self.kv_connector is not None and self.kv_connector != "TritonNixlConnector" and self.kv_role is None: ++ if self.kv_connector is not None and self.kv_connector != "DynemoNixlConnector" and self.kv_role is None: raise ValueError("Please specify kv_disagg_role when kv_connector " "is set, supported roles are `kv_producer`, " "`kv_consumer`, and `kv_both`") @@ -44,7 +44,7 @@ index 9ba49757..a2f88854 100644 def need_kv_parallel_group(self) -> bool: # for those database-based connector, vLLM does not need to create # parallel group, and in that case the kv parallel size will be 1. -+ if self.kv_connector == "TritonNixlConnector": ++ if self.kv_connector == "DynemoNixlConnector": + return False return self.kv_connector is not None and self.kv_parallel_size > 1 @@ -277,7 +277,7 @@ index 00000000..350453cd +logger = logging.getLogger(__name__) + + -+class TritonResult: ++class DynemoResult: + OK = 0 + ERR = 1 + @@ -290,12 +290,12 @@ index 00000000..350453cd + + try: + self.lib = ctypes.CDLL(lib_path) -+ self.lib.triton_llm_init.argtypes = [c_char_p, c_char_p, c_int64] -+ self.lib.triton_llm_init.restype = c_uint32 ++ self.lib.dynemo_llm_init.argtypes = [c_char_p, c_char_p, c_int64] ++ self.lib.dynemo_llm_init.restype = c_uint32 + -+ result = self.lib.triton_llm_init(namespace.encode(), ++ result = self.lib.dynemo_llm_init(namespace.encode(), + component.encode(), worker_id) -+ if result == TritonResult.OK: ++ if result == DynemoResult.OK: + logger.info( + "KVCacheEventManager initialized successfully. Ready to publish KV Cache Events" + ) @@ -306,7 +306,7 @@ index 00000000..350453cd + print(f"Failed to load {lib_path}") + raise e + -+ self.lib.triton_kv_event_publish_stored.argtypes = [ ++ self.lib.dynemo_kv_event_publish_stored.argtypes = [ + ctypes.c_uint64, # event_id + ctypes.POINTER(ctypes.c_uint32), # token_ids + ctypes.POINTER(ctypes.c_size_t), # num_block_tokens @@ -315,14 +315,14 @@ index 00000000..350453cd + ctypes.POINTER(ctypes.c_uint64), # parent_hash + ctypes.c_uint64, # lora_id + ] -+ self.lib.triton_kv_event_publish_stored.restype = ctypes.c_uint32 # triton_llm_result_t ++ self.lib.dynemo_kv_event_publish_stored.restype = ctypes.c_uint32 # dynemo_llm_result_t + -+ self.lib.triton_kv_event_publish_removed.argtypes = [ ++ self.lib.dynemo_kv_event_publish_removed.argtypes = [ + ctypes.c_uint64, # event_id + ctypes.POINTER(ctypes.c_uint64), # block_ids + ctypes.c_size_t, # num_blocks + ] -+ self.lib.triton_kv_event_publish_removed.restype = ctypes.c_uint32 # triton_llm_result_t ++ self.lib.dynemo_kv_event_publish_removed.restype = ctypes.c_uint32 # dynemo_llm_result_t + + self.event_id_counter = 0 + @@ -336,7 +336,7 @@ index 00000000..350453cd + if parent is not None else None) + + # Publish the event -+ result = self.lib.triton_kv_event_publish_stored( ++ result = self.lib.dynemo_kv_event_publish_stored( + self.event_id_counter, # uint64_t event_id + token_ids_arr, # const uint32_t *token_ids + num_block_tokens, # const uintptr_t *num_block_tokens @@ -346,7 +346,7 @@ index 00000000..350453cd + 0, # uint64_t lora_id + ) + -+ if result == TritonResult.OK: ++ if result == DynemoResult.OK: + logger.debug(f"Store - Published KV Event: {block.content_hash}") + else: + logger.debug( @@ -355,28 +355,23 @@ index 00000000..350453cd + self.event_id_counter += 1 + + def enqueue_removed_event(self, block_hash: PrefixHash): -+ result = self.lib.triton_kv_event_publish_removed( ++ result = self.lib.dynemo_kv_event_publish_removed( + self.event_id_counter, + (ctypes.c_uint64 * 1)(block_hash), + 1, + ) + -+ if result == TritonResult.OK: ++ if result == DynemoResult.OK: + logger.debug(f"Remove - Published KV Event: {block_hash}") + else: + logger.debug(f"Remove - Failed to Publish KV Event: {block_hash}") + + self.event_id_counter += 1 diff --git a/vllm/core/scheduler.py b/vllm/core/scheduler.py -index f507847a..abe574d1 100644 +index f507847a..ee20d50c 100644 --- a/vllm/core/scheduler.py +++ b/vllm/core/scheduler.py -@@ -4,22 +4,22 @@ import enum - import os - import random - import time -+import copy - from collections import deque +@@ -8,18 +8,17 @@ from collections import deque from dataclasses import dataclass, field from typing import Callable, Deque, Dict, Iterable, List, Optional from typing import Sequence as GenericSequence @@ -398,7 +393,7 @@ index f507847a..abe574d1 100644 logger = init_logger(__name__) # Test-only. If configured, decode is preempted with -@@ -325,12 +325,14 @@ class Scheduler: +@@ -325,12 +324,14 @@ class Scheduler: def __init__( self, @@ -413,7 +408,7 @@ index f507847a..abe574d1 100644 self.scheduler_config = scheduler_config self.cache_config = cache_config # Note for LoRA scheduling: the current policy is extremely -@@ -356,6 +358,7 @@ class Scheduler: +@@ -356,6 +357,7 @@ class Scheduler: # Create the block space manager. self.block_manager = BlockSpaceManagerImpl( @@ -421,7 +416,7 @@ index f507847a..abe574d1 100644 block_size=self.cache_config.block_size, num_gpu_blocks=num_gpu_blocks, num_cpu_blocks=num_cpu_blocks, -@@ -371,6 +374,16 @@ class Scheduler: +@@ -371,6 +373,14 @@ class Scheduler: # Sequence groups in the SWAPPED state. # Contain decode requests that are swapped out. self.swapped: Deque[SequenceGroup] = deque() @@ -429,8 +424,6 @@ index f507847a..abe574d1 100644 + # Sequence groups in the REMOTE_PREFILLING state. + # Contain requests that are being prefilled by a remote worker. + self.remote_prefilling: Deque[SequenceGroup] = deque() -+ # Contain requests that are being prefilled by a local worker. -+ self.prefill_sending: Deque[SequenceGroup] = deque() + + self._remote_prefill_outputs: Dict[str, int] = {} + @@ -438,25 +431,24 @@ index f507847a..abe574d1 100644 # Sequence groups finished requests ids since last step iteration. # It lets the model know that any state associated with these requests # can and must be released after the current step. -@@ -501,7 +514,7 @@ class Scheduler: +@@ -501,7 +511,7 @@ class Scheduler: def has_unfinished_seqs(self) -> bool: return len(self.waiting) != 0 or len(self.running) != 0 or len( - self.swapped) != 0 -+ self.swapped) != 0 or len(self.remote_prefilling) != 0 or len(self.prefill_sending) != 0 ++ self.swapped) != 0 or len(self.remote_prefilling) != 0 def get_prefix_cache_hit_rate(self, device: Device) -> float: return self.block_manager.get_prefix_cache_hit_rate(device) -@@ -523,6 +536,8 @@ class Scheduler: +@@ -523,6 +533,7 @@ class Scheduler: budget: SchedulingBudget, curr_loras: Optional[Set[int]], enable_chunking: bool = False, -+ finished_prefills: Optional[Set[str]] = None, -+ finished_transfers: Optional[Set[str]] = None ++ finished_prefills: Optional[Set[str]] = None ) -> SchedulerRunningOutputs: """Schedule sequence groups that are running. -@@ -537,6 +552,8 @@ class Scheduler: +@@ -537,6 +548,8 @@ class Scheduler: chunked number of tokens are scheduled if `budget.num_batched_tokens` has not enough capacity to schedule all tokens. @@ -465,7 +457,7 @@ index f507847a..abe574d1 100644 Returns: SchedulerRunningOutputs. -@@ -566,6 +583,38 @@ class Scheduler: +@@ -566,6 +579,24 @@ class Scheduler: preempted: List[SequenceGroup] = ret.preempted swapped_out: List[SequenceGroup] = ret.swapped_out @@ -476,7 +468,6 @@ index f507847a..abe574d1 100644 + if seq_group.request_id not in finished_prefills: + leftover_remote_prefilling_sequences.append(seq_group) + continue -+ + else: + finished_prefills.remove(seq_group.request_id) + assert len(seq_group.seqs) == 1 @@ -487,63 +478,39 @@ index f507847a..abe574d1 100644 + seq.data._stage = SequenceStage.DECODE + self.running.appendleft(seq_group) + remote_prefilling_queue.extendleft(leftover_remote_prefilling_sequences) -+ -+ remote_transfers_queue = self.prefill_sending -+ leftover_remote_transfers_sequences: Deque[SequenceGroup] = deque() -+ while remote_transfers_queue: -+ seq_group = remote_transfers_queue.popleft() -+ if seq_group.request_id not in finished_transfers: -+ leftover_remote_transfers_sequences.append(seq_group) -+ else: -+ finished_transfers.remove(seq_group.request_id) -+ assert len(seq_group.seqs) == 1 -+ seq = seq_group.seqs[0] -+ self.free_seq(seq) -+ remote_transfers_queue.extendleft(leftover_remote_transfers_sequences) + running_queue = self.running assert len(self._async_stopped) == 0 while running_queue: -@@ -1008,7 +1057,17 @@ class Scheduler: +@@ -1008,7 +1039,7 @@ class Scheduler: if curr_loras is not None and lora_int_id > 0: curr_loras.add(lora_int_id) waiting_queue.popleft() - self._allocate_and_set_running(seq_group) -+ -+ seq_group_copy = copy.deepcopy(seq_group) -+ seq_group_copy.seqs[0].seq_id = seq_group.seqs[0].seq_id + 1 -+ -+ logger.debug("Allocating and setting running or remote prefill for seq_group %s", seq_group.request_id) -+ logger.debug("Seq id: %s", seq_group.seqs[0].seq_id) + self._allocate_and_set_running_or_remote_prefill(seq_group) -+ if seq_group.remote_prefill_params is not None and seq_group.remote_prefill_params.is_remote_decode: -+ logger.debug("Seq id: %s", seq_group_copy.seqs[0].seq_id) -+ self._allocate_and_set_running_or_remote_prefill(seq_group_copy) -+ self.prefill_sending.append(seq_group_copy) if enable_chunking and self.scheduler_config.is_multi_step: blocks_to_copy: List[Tuple[int, int]] = [] -@@ -1048,7 +1107,7 @@ class Scheduler: +@@ -1048,7 +1079,7 @@ class Scheduler: num_lookahead_slots=self._get_num_lookahead_slots( is_prefill=True, enable_chunking=enable_chunking)) - def _schedule_default(self) -> SchedulerOutputs: -+ def _schedule_default(self, finished_prefills: Optional[Set[str]] = None, finished_transfers: Optional[Set[str]] = None) -> SchedulerOutputs: ++ def _schedule_default(self, finished_prefills: Optional[Set[str]] = None) -> SchedulerOutputs: """Schedule queued requests. The current policy is designed to optimize the throughput. First, -@@ -1090,7 +1149,9 @@ class Scheduler: +@@ -1090,7 +1121,8 @@ class Scheduler: if len(prefills.seq_groups) == 0: running_scheduled = self._schedule_running(budget, curr_loras, - enable_chunking=False) + enable_chunking=False, -+ finished_prefills=finished_prefills, -+ finished_transfers=finished_transfers) ++ finished_prefills=finished_prefills) # If any sequence group is preempted, do not swap in any sequence # group. because it means there's no slot for new running requests. -@@ -1106,7 +1167,12 @@ class Scheduler: +@@ -1106,7 +1138,12 @@ class Scheduler: self.waiting.extendleft(running_scheduled.preempted) # Update new running requests. if len(prefills.seq_groups) > 0: @@ -557,31 +524,30 @@ index f507847a..abe574d1 100644 self.running.extend(running_scheduled.decode_seq_groups_list) -@@ -1248,12 +1314,14 @@ class Scheduler: +@@ -1248,12 +1285,14 @@ class Scheduler: len(running_scheduled.swapped_out)), ) - def _schedule(self) -> SchedulerOutputs: -+ def _schedule(self, finished_prefills: Optional[Set[str]] = None, finished_transfers: Optional[Set[str]] = None) -> SchedulerOutputs: ++ def _schedule(self, finished_prefills: Optional[Set[str]] = None) -> SchedulerOutputs: """Schedule queued requests.""" if self.scheduler_config.chunked_prefill_enabled: -+ if finished_prefills or finished_transfers: ++ if finished_prefills: + raise ValueError("Chunked prefill does not support remote prefills") return self._schedule_chunked_prefill() else: - return self._schedule_default() -+ return self._schedule_default(finished_prefills, finished_transfers) ++ return self._schedule_default(finished_prefills) def _can_append_slots(self, seq_group: SequenceGroup, enable_chunking: bool) -> bool: -@@ -1287,14 +1355,16 @@ class Scheduler: +@@ -1287,14 +1326,15 @@ class Scheduler: return no_single_seq def schedule( - self + self, -+ finished_prefills: Optional[Set[str]] = None, -+ finished_transfers: Optional[Set[str]] = None ++ finished_prefills: Optional[Set[str]] = None ) -> Tuple[List[SequenceGroupMetadata], SchedulerOutputs, bool]: # Schedule sequence groups. # This function call changes the internal states of the scheduler @@ -590,11 +556,11 @@ index f507847a..abe574d1 100644 - scheduler_outputs: SchedulerOutputs = self._schedule() + scheduler_start_time = time.perf_counter() -+ scheduler_outputs: SchedulerOutputs = self._schedule(finished_prefills, finished_transfers) ++ scheduler_outputs: SchedulerOutputs = self._schedule(finished_prefills) now = time.time() if not self.cache_config.enable_prefix_caching: -@@ -1333,7 +1403,8 @@ class Scheduler: +@@ -1333,7 +1373,8 @@ class Scheduler: encoder_seq_data = None cross_block_table = None @@ -604,24 +570,18 @@ index f507847a..abe574d1 100644 seq_id = seq.seq_id seq_data[seq_id] = seq.data block_tables[seq_id] = self.block_manager.get_block_table(seq) -@@ -1364,9 +1435,16 @@ class Scheduler: +@@ -1364,6 +1405,10 @@ class Scheduler: < seqs[0].data.get_len()): do_sample = False + is_remote_prefill = False + if is_first_prefill and seq_group.remote_prefill_params is not None and seq_group.remote_prefill_params.is_remote_prefill: + is_remote_prefill = True -+ if is_first_prefill and seq_group.remote_prefill_params is not None and seq_group.remote_prefill_params.is_remote_decode: -+ block_tables[seq_group.seqs[0].seq_id + 1] = self.block_manager.block_tables[seq.seq_id + 1].physical_block_ids + # It assumes the scheduled_seq_groups is ordered by # prefill < decoding. if is_first_prefill or not self.scheduler_config.send_delta_data: -+ logger.debug("Assinged blocks: %s", block_tables) - seq_group_metadata = SequenceGroupMetadata( - request_id=seq_group.request_id, - is_prompt=is_prompt, -@@ -1392,6 +1470,7 @@ class Scheduler: +@@ -1392,6 +1437,7 @@ class Scheduler: if scheduler_outputs.num_prefill_groups > 0 else None, mm_processor_kwargs=seq_group.mm_processor_kwargs, prompt_adapter_request=seq_group.prompt_adapter_request, @@ -629,7 +589,7 @@ index f507847a..abe574d1 100644 ) else: # When SPMD mode is enabled, we only send delta data except for -@@ -1490,10 +1569,13 @@ class Scheduler: +@@ -1490,10 +1536,13 @@ class Scheduler: self._async_stopped.clear() @@ -645,80 +605,12 @@ index f507847a..abe574d1 100644 def _append_slots(self, seq_group: SequenceGroup, -diff --git a/vllm/distributed/device_communicators/kv_rearrange.py b/vllm/distributed/device_communicators/kv_rearrange.py -new file mode 100644 -index 00000000..9b938039 ---- /dev/null -+++ b/vllm/distributed/device_communicators/kv_rearrange.py -@@ -0,0 +1,61 @@ -+import torch -+import triton -+import triton.language as tl -+ -+@triton.jit -+def rearrange_kernel( -+ t1_ptr, -+ t2_ptr, -+ N, -+ B, -+ H, -+ C, -+ d, -+ tensor_subset_size, -+ block_size, -+ token_size, -+ BLOCK_SIZE: tl.constexpr, -+): -+ pid = tl.program_id(0) -+ -+ block_start = pid * BLOCK_SIZE -+ offsets = block_start + tl.arange(0, BLOCK_SIZE) -+ -+ curr_n = offsets // block_size -+ curr_b = offsets // token_size % B -+ curr_h = offsets // C % H -+ curr_c = offsets % C -+ -+ src_pos = offsets -+ -+ tp_group = curr_h * d // H -+ dst_h = curr_h % (H // d) -+ tp_group_offset = curr_n * (block_size // d) + curr_b * (H // d) * C + dst_h * C + curr_c -+ -+ dst_pos = tensor_subset_size * tp_group + tp_group_offset -+ -+ tl.store(t2_ptr + dst_pos, tl.load(t1_ptr + src_pos)) -+ -+def rearrange_tensors(t1: torch.Tensor, t2: torch.Tensor, d: int): -+ N, B, H, C = t1.shape -+ -+ assert t2.shape == (N, B, H, C), "Destination tensor must have same shape as source" -+ assert H % d == 0, "H must be divisible by d" -+ -+ block_size = B * H * C -+ token_size = H * C -+ tensor_size = N * block_size -+ tensor_subset_size = tensor_size // d -+ -+ BLOCK_SIZE = 1024 -+ grid = ((N * B * H * C + BLOCK_SIZE - 1) // BLOCK_SIZE,) -+ -+ rearrange_kernel[grid]( -+ t1, t2, -+ N, B, H, C, -+ d, -+ tensor_subset_size, -+ block_size, -+ token_size, -+ BLOCK_SIZE=BLOCK_SIZE -+ ) -\ No newline at end of file diff --git a/vllm/distributed/device_communicators/nixl.py b/vllm/distributed/device_communicators/nixl.py new file mode 100644 -index 00000000..f1618bc4 +index 00000000..bc962726 --- /dev/null +++ b/vllm/distributed/device_communicators/nixl.py -@@ -0,0 +1,318 @@ +@@ -0,0 +1,249 @@ +import torch +from typing import List, Tuple +from vllm.config import VllmConfig @@ -726,18 +618,39 @@ index 00000000..f1618bc4 +import msgspec +import time +import uuid -+from collections import defaultdict -+from .kv_rearrange import rearrange_tensors ++from nixl_wrapper import nixl_wrapper as NixlWrapper + +logger = init_logger(__name__) + -+# Lazy import nixl_wrapper to avoid loading nixl_bindings if nixl is not used -+try: -+ from nixl_wrapper import nixl_wrapper as NixlWrapper # type: ignore -+ logger.info("NIXL is available") -+except ImportError: -+ logger.warning("NIXL is not available") -+ NixlWrapper = None ++ ++def nixl_wrapper_init_patch(self, agent_name, nixl_config): ++ logger.info("Initializing patched NixlWrapper") ++ import nixl_bindings as nixl ++ # Read available backends and device info from nixl_config ++ # For now setting the multithreading to enabled. ++ devices = nixl.nixlAgentConfig(False) ++ init = nixl.nixlUcxInitParams() ++ ++ self.name = agent_name ++ self.notifs = {} ++ self.backends = {} ++ self.agent = nixl.nixlAgent(agent_name, devices) ++ self.backends["UCX"] = self.agent.createBackend(init) ++ ++ self.nixl_mems = {"DRAM": nixl.DRAM_SEG, ++ "VRAM": nixl.VRAM_SEG, ++ "cpu": nixl.DRAM_SEG, ++ "cuda": nixl.VRAM_SEG} ++ self.nixl_ops = {"WRITE": nixl.NIXL_WR_FLUSH, ++ "READ": nixl.NIXL_RD_FLUSH, ++ "WRITE_NOTIF": nixl.NIXL_WR_NOTIF, ++ "READ_NOTIF": nixl.NIXL_RD_NOTIF} ++ ++ print("Initializied NIXL agent:", agent_name) ++ ++NixlWrapper.__init__ = nixl_wrapper_init_patch ++ ++ + +class NixlMetadata( + msgspec.Struct, @@ -749,20 +662,14 @@ index 00000000..f1618bc4 + kv_caches_base_addr: List[List[Tuple[int, int]]] # base address for each rank for each layer for keys and values + + -+class TritonNixlConnector: ++class DynemoNixlConnector: + def __init__(self, vllm_config: VllmConfig, engine_id: str, rank: int): + self.vllm_config = vllm_config -+ if NixlWrapper is None: -+ logger.error("NIXL is not available") -+ raise RuntimeError("NIXL is not available") -+ logger.info("Initializing NIXL wrapper") + self.nixl_wrapper = NixlWrapper(str(uuid.uuid4()), None) + + self.num_layers = None + self.num_blocks = None -+ self.num_heads = None + self.block_len = None -+ self.kv_caches = None + self.kv_caches_base_addr = {} + self.kv_cache_shape = {} + @@ -771,51 +678,33 @@ index 00000000..f1618bc4 + self.engine_id = engine_id + self.rank = rank + self.notifs = {} -+ self._tp_size = {} -+ self._block_descs = {} -+ self._xfer_side_handles = {} -+ -+ -+ self._transfers = defaultdict(list) -+ -+ -+ self._tp_size[engine_id] = vllm_config.parallel_config.tensor_parallel_size -+ + + @property + def agent_name(self): + return self.nixl_wrapper.name + + def register_kv_caches(self, kv_caches: List[torch.Tensor]): -+ _, num_blocks, block_size, num_heads, head_dim = kv_caches[0].shape ++ caches_data = [] ++ self.num_layers = len(kv_caches) ++ _, _, block_size, num_heads, head_dim = kv_caches[0].shape + self.block_len = block_size * num_heads * head_dim * kv_caches[0].element_size() + logger.debug("Per layer kv cache size: %s", kv_caches[0].shape) -+ self.num_layers = len(kv_caches) -+ self.num_blocks = num_blocks -+ self.num_heads = num_heads -+ self.kv_caches = kv_caches ++ + kv_caches_base_addr = [] -+ caches_data = [] -+ blocks_data = [] + for key_cache, value_cache in kv_caches: + for cache in [key_cache, value_cache]: + base_addr = cache.data_ptr() -+ region_len = num_blocks * self.block_len -+ caches_data.append((base_addr, region_len, self.rank)) -+ for block_id in range(self.num_blocks): -+ blocks_data.append((base_addr + block_id * self.block_len, self.block_len, self.rank)) -+ ++ region_len = cache.numel() * cache.element_size() ++ gpu_id = cache.get_device() ++ assert gpu_id > -1, "Tensor is not on GPU" ++ caches_data.append((base_addr, region_len, gpu_id)) + kv_caches_base_addr.append((key_cache.data_ptr(), value_cache.data_ptr())) + self.kv_caches_base_addr[self.engine_id] = kv_caches_base_addr + + descs = self.nixl_wrapper.get_descs(("VRAM", caches_data)) -+ logger.debug("Registering descs: %s", caches_data) + self.nixl_wrapper.register_memory(descs) + self._registered_descs.append(descs) + -+ self._block_descs[self.engine_id] = self.nixl_wrapper.get_descs(("VRAM", blocks_data)) -+ self._xfer_side_handles[self.engine_id] = self.nixl_wrapper.prep_xfer_side(self._block_descs[self.engine_id]) -+ + def get_agent_metadata(self): + return self.nixl_wrapper.get_agent_metadata() + @@ -825,14 +714,10 @@ index 00000000..f1618bc4 + for agent_name in self._remote_agents.values(): + self.nixl_wrapper.remove_remote_agent(agent_name) + -+ def add_remote_agent(self, engine_id, agent_metadata, agent_tp): -+ self._tp_size[engine_id] = agent_tp -+ agent_names = [] -+ for agent_meta in agent_metadata: -+ agent_name = self.nixl_wrapper.add_remote_agent(agent_meta) -+ agent_names.append(agent_name) -+ self._remote_agents[engine_id] = agent_names -+ return agent_names ++ def add_remote_agent(self, engine_id, agent_metadata): ++ agent_name = self.nixl_wrapper.add_remote_agent(agent_metadata) ++ self._remote_agents[engine_id] = agent_name ++ return agent_name + + def get_descs_ids(self, layer_ids, block_ids): + if layer_ids == "all": @@ -847,29 +732,17 @@ index 00000000..f1618bc4 + descs_ids.append(2 * (self.num_blocks * layer_id + block_id) + 1) + return descs_ids + -+ def _get_range_descs(self, ranges, layer_ids, kv_caches_base_addr, tp_multiplier=1, rank=None, i=0): -+ if rank is None: -+ rank = self.rank -+ offset_block_len = self.block_len -+ block_len = self.block_len // tp_multiplier -+ tp_offset = i * block_len -+ else: -+ offset_block_len = self.block_len // tp_multiplier -+ block_len = self.block_len // tp_multiplier -+ tp_offset = 0 -+ logger.debug("Getting range descs for layer ids: %s, ranges: %s, tp_multiplier: %s, rank: %s, i: %s", layer_ids, ranges, tp_multiplier, rank, i) ++ def _get_range_descs(self, engine_id, ranges, layer_ids): + if layer_ids == "all": + layer_ids = list(range(self.num_layers)) + blocks_data = [] + for layer_id in layer_ids: + for range_start, range_end in ranges: -+ range_len = range_end - range_start + 1 -+ key_base_addr, value_base_addr = kv_caches_base_addr[layer_id] -+ start_offset = range_start * offset_block_len + tp_offset * range_len -+ blocks_len = range_len * block_len -+ blocks_data.append((key_base_addr + start_offset, blocks_len, rank)) -+ blocks_data.append((value_base_addr + start_offset, blocks_len, rank)) -+ logger.debug("Blocks data: %s", blocks_data) ++ key_base_addr, value_base_addr = self.kv_caches_base_addr[engine_id][layer_id] ++ start_offset = range_start * self.block_len ++ blocks_len = (range_end - range_start + 1) * self.block_len ++ blocks_data.append((key_base_addr + start_offset, blocks_len, self.rank)) ++ blocks_data.append((value_base_addr + start_offset, blocks_len, self.rank)) + return self.nixl_wrapper.get_descs(("VRAM", blocks_data)) + + def _get_ranges(self, block_ids): @@ -882,9 +755,9 @@ index 00000000..f1618bc4 + ranges = [] + for i in range(len(sorted_block_ids)): + if i == 0 or sorted_block_ids[i] != sorted_block_ids[i-1] + 1: -+ ranges.append([sorted_block_ids[i], sorted_block_ids[i]]) ++ ranges.append([sorted_block_ids[i]]) + else: -+ ranges[-1][1] = sorted_block_ids[i] ++ ranges[-1].append(sorted_block_ids[i]) + return ranges + + def _get_same_length_ranges(self, src_ranges, dst_ranges): @@ -924,24 +797,11 @@ index 00000000..f1618bc4 + src_idx += 1 + + return src_overlapping_ranges, dst_overlapping_ranges -+ -+ -+ -+ def _get_block_descs_ids(self, layer_ids, block_ids): -+ if layer_ids == "all": -+ layer_ids = list(range(self.num_layers)) -+ if block_ids == "all": -+ block_ids = list(range(self.num_blocks)) -+ descs_ids = [] -+ for layer_id in layer_ids: -+ for is_value in [0, 1]: -+ for block_id in block_ids: -+ descs_ids.append(layer_id * 2 * self.num_blocks + is_value * self.num_blocks + block_id) -+ return descs_ids + + + -+ def transfer_mem(self, src_block_ids, staging_block_ids, dst_block_ids, dst_engine_id, notify_msg, use_prepped_xfer=False): ++ def transfer_mem(self, src_block_ids, dst_block_ids, dst_engine_id, notify_msg): ++ + start_time = time.perf_counter() + logger.debug("Transferring memory from %s to %s with notify message %s", self.agent_name, dst_engine_id, notify_msg) + @@ -950,62 +810,44 @@ index 00000000..f1618bc4 + # If isl equals to a multiple of tokens_per_block + 1, prefill engine will have \ + # one less block due to the missing last token. + dst_block_ids = dst_block_ids[:len(src_block_ids)] -+ assert len(staging_block_ids) == len(src_block_ids) + -+ if use_prepped_xfer: -+ raise NotImplementedError("Prepped xfer is not implemented") -+ # src_block_descs_ids = self._get_block_descs_ids("all", src_block_ids) -+ # dst_block_descs_ids = self._get_block_descs_ids("all", dst_block_ids) -+ -+ # src_xfer_side_handle = self._xfer_side_handles[self.engine_id] -+ # dst_xfer_side_handle = self._xfer_side_handles[dst_engine_id] -+ -+ # logger.debug("Time to get block desc ids: %s ms", (time.perf_counter() - start_time) * 1000) -+ -+ # handle = self.nixl_wrapper.make_prepped_xfer(src_xfer_side_handle, src_block_descs_ids, -+ # dst_xfer_side_handle, dst_block_descs_ids, -+ # notify_msg, "WRITE", no_check=True) -+ # else: -+ # Legacy path using range-based transfers + src_ranges = self._get_ranges(src_block_ids) -+ staging_ranges = self._get_ranges(staging_block_ids) + dst_ranges = self._get_ranges(dst_block_ids) ++ src_overlapping_ranges, dst_overlapping_ranges = self._get_same_length_ranges(src_ranges, dst_ranges) + -+ assert len(src_ranges) == 1 -+ assert len(staging_ranges) == 1 ++ logger.debug("Got %s overlapping ranges for %s blocks", len(src_overlapping_ranges), len(src_block_ids)) + -+ tp_multiplier = self._tp_size[dst_engine_id] // self._tp_size[self.engine_id] -+ -+ src_range_start, src_range_end = src_ranges[0] -+ src_range_len = src_range_end - src_range_start + 1 -+ staging_range_start, staging_range_end = staging_ranges[0] -+ staging_range_len = staging_range_end - staging_range_start + 1 ++ logger.debug("Time to get ranges: %s ms", time.perf_counter() - start_time) + -+ logger.debug("Rearranging tensors for cache: %s, src_ranges: %s of len %s, staging_ranges: %s of len %s", self.kv_caches[0].shape, src_ranges, src_range_len, staging_ranges, staging_range_len) -+ for kv_cache in self.kv_caches: -+ for cache in kv_cache: -+ rearrange_tensors(cache[src_range_start:src_range_start + src_range_len], cache[staging_range_start:staging_range_start + staging_range_len], tp_multiplier) ++ src_descs = self._get_range_descs(self.engine_id, src_overlapping_ranges, "all") ++ dst_descs = self._get_range_descs(dst_engine_id, dst_overlapping_ranges, "all") + -+ staging_overlapping_ranges, dst_overlapping_ranges = self._get_same_length_ranges(staging_ranges, dst_ranges) -+ assert len(staging_overlapping_ranges) == len(dst_overlapping_ranges) ++ logger.debug("Time to get descs: %s ms", (time.perf_counter() - start_time) * 1000) + -+ for i in range(tp_multiplier): ++ handle = self.nixl_wrapper.initialize_xfer(src_descs, dst_descs, self._remote_agents[dst_engine_id], notify_msg, "WRITE") ++ logger.debug("Time to initialize xfer: %s ms", (time.perf_counter() - start_time) * 1000) ++ logger.debug("Transfer handle: %s", handle) ++ status = self.nixl_wrapper.transfer(handle) ++ logger.debug("Time to transfer: %s ms", (time.perf_counter() - start_time) * 1000) ++ logger.debug("Transfer status: %s", status) ++ # TODO ptarasiewicz: remove blocking transfer mem ++ # add scheduler check for transfer done ++ while True: ++ xfer_state = self.nixl_wrapper.check_xfer_state(handle) ++ if xfer_state == "ERR": ++ raise RuntimeError("Transfer failed") ++ elif xfer_state == "DONE": ++ logger.debug("Transfer done") ++ break ++ elif xfer_state == "PROC": ++ time.sleep(0.01) ++ else: ++ raise RuntimeError("Unknown transfer state") ++ logger.debug("Time to wait for transfer: %s ms", (time.perf_counter() - start_time) * 1000) ++ self.nixl_wrapper.abort_xfer(handle) ++ logger.debug("Time to abort xfer: %s ms", (time.perf_counter() - start_time) * 1000) ++ logger.debug("Transfer time: %s ms", (time.perf_counter() - start_time) * 1000) + -+ src_descs = self._get_range_descs(staging_overlapping_ranges, "all", self.kv_caches_base_addr[self.engine_id], tp_multiplier, i=i) -+ dst_descs = self._get_range_descs(dst_overlapping_ranges, "all", self.kv_caches_base_addr[dst_engine_id][self.rank * tp_multiplier + i], tp_multiplier, rank=self.rank * tp_multiplier + i) -+ logger.debug("Time to get descs: %s ms", (time.perf_counter() - start_time) * 1000) -+ -+ logger.debug("Transfering to agent %s", self._remote_agents[dst_engine_id][self.rank * tp_multiplier + i]) -+ handle = self.nixl_wrapper.initialize_xfer(src_descs, dst_descs, -+ self._remote_agents[dst_engine_id][self.rank * tp_multiplier + i], -+ notify_msg, "WRITE") -+ self._transfers[notify_msg].append(handle) -+ logger.debug("Time to initialize xfer: %s ms", (time.perf_counter() - start_time) * 1000) -+ logger.debug("Transfer handle: %s", handle) -+ status = self.nixl_wrapper.transfer(handle) -+ logger.debug("Time to transfer: %s ms", (time.perf_counter() - start_time) * 1000) -+ logger.debug("Transfer status: %s", status) -+ + def deserialize_descs(self, serialized_descs): + return self.nixl_wrapper.deserialize_descs(serialized_descs) + @@ -1018,26 +860,6 @@ index 00000000..f1618bc4 + + def add_remote_kv_caches_base_addr(self, engine_id, kv_caches_base_addr): + self.kv_caches_base_addr[engine_id] = kv_caches_base_addr -+ -+ def get_done_tranfers(self) -> List[str]: -+ done_req_ids = [] -+ for req_id, handles in self._transfers.items(): -+ running_reqs = [] -+ for handle in handles: -+ xfer_state = self.nixl_wrapper.check_xfer_state(handle) -+ if xfer_state == "DONE": -+ # self.nixl_wrapper.abort_xfer(handle) # TODO ptarasiewicz: why abort is throwing errors? -+ continue -+ if xfer_state == "PROC": -+ running_reqs.append(handle) -+ else: -+ raise RuntimeError("Transfer failed with state %s", xfer_state) -+ if len(running_reqs) == 0: -+ done_req_ids.append(req_id) -+ else: -+ self._transfers[req_id] = running_reqs -+ return done_req_ids -\ No newline at end of file diff --git a/vllm/distributed/kv_transfer/kv_connector/factory.py b/vllm/distributed/kv_transfer/kv_connector/factory.py index fe480533..61a357d0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/factory.py @@ -1064,9 +886,9 @@ index fe480533..61a357d0 100644 "SimpleConnector") + +KVConnectorFactory.register_connector( -+ "TritonNcclConnector", -+ "vllm.distributed.kv_transfer.kv_connector.triton_connector", -+ "TritonConnector") ++ "DynemoNcclConnector", ++ "vllm.distributed.kv_transfer.kv_connector.dynemo_connector", ++ "DynemoConnector") \ No newline at end of file diff --git a/vllm/distributed/kv_transfer/kv_connector/simple_connector.py b/vllm/distributed/kv_transfer/kv_connector/simple_connector.py index 2033e976..e33919c1 100644 @@ -1396,7 +1218,7 @@ index 2033e976..e33919c1 100644 + world_group.broadcast_object(kv_config_enhanced) + + else: -+ raise NotImplementedError("MooncakeConnector is not supported in Triton Distributed vllm patch") ++ raise NotImplementedError("MooncakeConnector is not supported in Dynemo Distributed vllm patch") + else: + kv_config_enhanced = world_group.broadcast_object() + logger.info("kv_config_enhanced: %s", kv_config_enhanced) @@ -1407,11 +1229,11 @@ index 2033e976..e33919c1 100644 + self.config.kv_consumers_pipeline_parallel_size = kv_config_enhanced["kv_consumers_pipeline_parallel_size"] + self.config.kv_producers_parallel_size = kv_config_enhanced["kv_producers_parallel_size"] \ No newline at end of file -diff --git a/vllm/distributed/kv_transfer/kv_connector/triton_connector.py b/vllm/distributed/kv_transfer/kv_connector/triton_connector.py +diff --git a/vllm/distributed/kv_transfer/kv_connector/dynemo_connector.py b/vllm/distributed/kv_transfer/kv_connector/dynemo_connector.py new file mode 100644 index 00000000..cb3b3660 --- /dev/null -+++ b/vllm/distributed/kv_transfer/kv_connector/triton_connector.py ++++ b/vllm/distributed/kv_transfer/kv_connector/dynemo_connector.py @@ -0,0 +1,350 @@ +# SPDX-License-Identifier: Apache-2.0 +""" @@ -1443,7 +1265,7 @@ index 00000000..cb3b3660 +logger = init_logger(__name__) + + -+class TritonConnector(KVConnectorBase): ++class DynemoConnector(KVConnectorBase): + + def __init__( + self, @@ -1457,16 +1279,16 @@ index 00000000..cb3b3660 + self.tp_size = config.parallel_config.tensor_parallel_size + self.rank = rank + -+ if self.config.kv_connector != "TritonNcclConnector": -+ raise NotImplementedError("Only TritonNcclConnector is supported by the TritonConnector class") ++ if self.config.kv_connector != "DynemoNcclConnector": ++ raise NotImplementedError("Only DynemoNcclConnector is supported by the DynemoConnector class") + + from vllm.distributed.kv_transfer.kv_pipe.pynccl_pipe import ( + PyNcclPipe) -+ from vllm.distributed.kv_transfer.kv_pipe.triton_nccl_pipe import ( -+ TritonNcclDataPlane) ++ from vllm.distributed.kv_transfer.kv_pipe.dynemo_nccl_pipe import ( ++ DynemoNcclDataPlane) + + logger.info( -+ "Initializing TritonNcclConnector under kv_transfer_config %s", ++ "Initializing DynemoNcclConnector under kv_transfer_config %s", + self.config) + + self.lookup_buffer_size = self.config.kv_buffer_size @@ -1498,7 +1320,7 @@ index 00000000..cb3b3660 + port_offset=port_offset_base, + ) + -+ self.data_plane = TritonNcclDataPlane( ++ self.data_plane = DynemoNcclDataPlane( + data_pipe=self.data_pipe, + port=self._get_data_plane_port(self.global_kv_rank), + ) @@ -2233,11 +2055,11 @@ index 7aa53d07..f5dd50b7 100644 def close(self): """ -diff --git a/vllm/distributed/kv_transfer/kv_pipe/triton_nccl_pipe.py b/vllm/distributed/kv_transfer/kv_pipe/triton_nccl_pipe.py +diff --git a/vllm/distributed/kv_transfer/kv_pipe/dynemo_nccl_pipe.py b/vllm/distributed/kv_transfer/kv_pipe/dynemo_nccl_pipe.py new file mode 100644 index 00000000..8a356504 --- /dev/null -+++ b/vllm/distributed/kv_transfer/kv_pipe/triton_nccl_pipe.py ++++ b/vllm/distributed/kv_transfer/kv_pipe/dynemo_nccl_pipe.py @@ -0,0 +1,124 @@ +import logging +import threading @@ -2253,7 +2075,7 @@ index 00000000..8a356504 +logger = logging.getLogger(__name__) + + -+class TritonNcclDataPlane: ++class DynemoNcclDataPlane: + def __init__( + self, + data_pipe: PyNcclPipe, @@ -2399,7 +2221,7 @@ index 321902d1..b8937ef8 100644 def ensure_model_parallel_initialized( diff --git a/vllm/engine/llm_engine.py b/vllm/engine/llm_engine.py -index d82d9ad9..254337cb 100644 +index d82d9ad9..9ba1a326 100644 --- a/vllm/engine/llm_engine.py +++ b/vllm/engine/llm_engine.py @@ -2,13 +2,17 @@ @@ -2482,13 +2304,11 @@ index d82d9ad9..254337cb 100644 + self.engine_id = str(uuid.uuid4()) + self._nixl_agents_names: Optional[List[str]] = None -+ if self.vllm_config.kv_transfer_config is not None and self.vllm_config.kv_transfer_config.kv_connector == "TritonNixlConnector": ++ if self.vllm_config.kv_transfer_config is not None and self.vllm_config.kv_transfer_config.kv_connector == "DynemoNixlConnector": + self._nixl_agents_names = self._initialize_nixl() + + self._request_notif_counter = defaultdict(lambda: -self.parallel_config.tensor_parallel_size) -+ self._request_done_counter = defaultdict(lambda: -self.parallel_config.tensor_parallel_size) + self._finished_prefills = set() -+ self._finished_transfers = set() + + @property + def is_nixl_initialized(self) -> bool: @@ -2507,6 +2327,8 @@ index d82d9ad9..254337cb 100644 + engine_id = nixl_metadata.engine_id + agents_metadata = nixl_metadata.agent_metadata + kv_caches_base_addr = nixl_metadata.kv_caches_base_addr ++ if len(agents_metadata) != len(self._nixl_agents_names): ++ raise ValueError("Number of agents does not match. Make sure all engines are initialized with the same parallel sizes.") + return self.model_executor.collective_rpc("add_remote_nixl_metadata", args=(engine_id, agents_metadata, kv_caches_base_addr)) + + def _initialize_nixl(self) -> List[bytes]: @@ -2540,16 +2362,7 @@ index d82d9ad9..254337cb 100644 ParallelSampleSequenceGroup.add_request( request_id, self, -@@ -574,6 +624,8 @@ class LLMEngine: - # Create the sequences. - block_size = self.cache_config.block_size - seq_id = next(self.seq_counter) -+ if remote_prefill_params is not None and remote_prefill_params.is_remote_decode: -+ next(self.seq_counter) # empty sequence for staging - eos_token_id = self.input_preprocessor.get_eos_token_id(lora_request) - - if is_encoder_decoder_inputs(processed_inputs): -@@ -584,7 +636,7 @@ class LLMEngine: +@@ -584,7 +634,7 @@ class LLMEngine: encoder_inputs = None seq = Sequence(seq_id, decoder_inputs, block_size, eos_token_id, @@ -2558,7 +2371,7 @@ index d82d9ad9..254337cb 100644 encoder_seq = (None if encoder_inputs is None else Sequence( seq_id, encoder_inputs, block_size, eos_token_id, lora_request, -@@ -601,8 +653,12 @@ class LLMEngine: +@@ -601,8 +651,12 @@ class LLMEngine: trace_headers=trace_headers, prompt_adapter_request=prompt_adapter_request, encoder_seq=encoder_seq, @@ -2572,7 +2385,7 @@ index d82d9ad9..254337cb 100644 seq_group = self._create_sequence_group_with_pooling( request_id, seq, -@@ -673,6 +729,7 @@ class LLMEngine: +@@ -673,6 +727,7 @@ class LLMEngine: trace_headers: Optional[Mapping[str, str]] = None, prompt_adapter_request: Optional[PromptAdapterRequest] = None, priority: int = 0, @@ -2580,7 +2393,7 @@ index d82d9ad9..254337cb 100644 *, inputs: Optional[PromptType] = None, # DEPRECATED ) -> None: -@@ -765,6 +822,7 @@ class LLMEngine: +@@ -765,6 +820,7 @@ class LLMEngine: prompt_adapter_request=prompt_adapter_request, trace_headers=trace_headers, priority=priority, @@ -2588,7 +2401,7 @@ index d82d9ad9..254337cb 100644 ) def _validate_token_prompt(self, prompt: PromptType, -@@ -799,6 +857,7 @@ class LLMEngine: +@@ -799,6 +855,7 @@ class LLMEngine: prompt_adapter_request: Optional[PromptAdapterRequest] = None, encoder_seq: Optional[Sequence] = None, priority: int = 0, @@ -2596,7 +2409,7 @@ index d82d9ad9..254337cb 100644 ) -> SequenceGroup: """Creates a SequenceGroup with SamplingParams.""" max_logprobs = self.get_model_config().max_logprobs -@@ -829,7 +888,9 @@ class LLMEngine: +@@ -829,7 +886,9 @@ class LLMEngine: trace_headers=trace_headers, prompt_adapter_request=prompt_adapter_request, encoder_seq=encoder_seq, @@ -2607,7 +2420,7 @@ index d82d9ad9..254337cb 100644 return seq_group -@@ -995,11 +1056,11 @@ class LLMEngine: +@@ -995,11 +1054,11 @@ class LLMEngine: # When we process only one request, no pop is required # (since later we will process all of the rest) (outputs, seq_group_metadata_list, scheduler_outputs, is_async, @@ -2621,7 +2434,7 @@ index d82d9ad9..254337cb 100644 # Sanity check assert len(seq_group_metadata_list) == len( -@@ -1325,15 +1386,49 @@ class LLMEngine: +@@ -1325,15 +1384,49 @@ class LLMEngine: # Clear outputs for each new scheduler iteration ctx.request_outputs.clear() @@ -2641,7 +2454,7 @@ index d82d9ad9..254337cb 100644 (seq_group_metadata_list, scheduler_outputs, allow_async_output_proc - ) = self.scheduler[virtual_engine].schedule() -+ ) = self.scheduler[virtual_engine].schedule(self._finished_prefills, self._finished_transfers) ++ ) = self.scheduler[virtual_engine].schedule(self._finished_prefills) + + + # Separate remote prefill and running seq groups @@ -2673,7 +2486,7 @@ index d82d9ad9..254337cb 100644 ctx.seq_group_metadata_list = seq_group_metadata_list ctx.scheduler_outputs = scheduler_outputs -@@ -1383,9 +1478,31 @@ class LLMEngine: +@@ -1383,9 +1476,29 @@ class LLMEngine: execute_model_req.async_callback = self.async_callbacks[ virtual_engine] @@ -2687,11 +2500,9 @@ index d82d9ad9..254337cb 100644 + req_id = scheduled_seq_group.seq_group.request_id + seq_id = scheduled_seq_group.seq_group.seqs[0].seq_id + block_table = seq_group_metadata.block_tables[seq_id] -+ staging_block_ids = seq_group_metadata.block_tables[seq_id + 1] + memory_transfer_req = MemoryTransferRequest( + request_id=req_id, + src_block_ids=block_table, -+ staging_block_ids=staging_block_ids, + dst_block_ids=remote_prefill_params.decode_block_ids, + dst_engine_id=remote_prefill_params.decode_engine_id, + notify_msg=req_id, @@ -2701,13 +2512,13 @@ index d82d9ad9..254337cb 100644 + + execute_model_req.memory_transfer_requests = memory_transfer_reqs + -+ outputs, request_notif_counter, request_done_counter = self.model_executor.execute_model( ++ outputs, request_notif_counter = self.model_executor.execute_model( execute_model_req=execute_model_req) - # We need to do this here so that last step's sampled_token_ids can # be passed to the next iteration for PP. if self.scheduler_config.is_multi_step: -@@ -1396,7 +1513,26 @@ class LLMEngine: +@@ -1396,7 +1509,20 @@ class LLMEngine: if len(ctx.output_queue) > 0: self._process_model_outputs(ctx=ctx) # No outputs in this case @@ -2718,7 +2529,7 @@ index d82d9ad9..254337cb 100644 + blocks_to_swap_out=[], + blocks_to_copy=[]) + -+ outputs, request_notif_counter, request_done_counter = self.model_executor.execute_model( ++ outputs, request_notif_counter = self.model_executor.execute_model( + execute_model_req=execute_model_req) + + for req_id, notif_count in request_notif_counter.items(): @@ -2726,16 +2537,10 @@ index d82d9ad9..254337cb 100644 + if self._request_notif_counter[req_id] > -1: + self._finished_prefills.add(req_id) + del self._request_notif_counter[req_id] -+ -+ for req_id, done_count in request_done_counter.items(): -+ self._request_done_counter[req_id] += done_count -+ if self._request_done_counter[req_id] > -1: -+ self._finished_transfers.add(req_id) -+ del self._request_done_counter[req_id] # Finish the current step for all the sequence groups. if self.scheduler_config.is_multi_step: -@@ -1456,7 +1592,7 @@ class LLMEngine: +@@ -1456,7 +1582,7 @@ class LLMEngine: # queued control plane messages, such as add/remove lora adapters. logger.debug("Stopping remote worker execution loop.") self.model_executor.stop_remote_worker_execution_loop() @@ -2813,7 +2618,7 @@ index 3cf1850e..6b90ece7 100644 + kv_active_blocks: int + kv_total_blocks: int diff --git a/vllm/engine/multiprocessing/client.py b/vllm/engine/multiprocessing/client.py -index 85b5f31e..c501e4c8 100644 +index 85b5f31e..d33d546a 100644 --- a/vllm/engine/multiprocessing/client.py +++ b/vllm/engine/multiprocessing/client.py @@ -8,6 +8,7 @@ from typing import (Any, AsyncGenerator, Dict, Iterator, List, Mapping, @@ -2895,7 +2700,7 @@ index 85b5f31e..c501e4c8 100644 + + @property + def using_nixl_connector(self) -> bool: -+ return self.vllm_config.kv_transfer_config is not None and self.vllm_config.kv_transfer_config.kv_connector == "TritonNixlConnector" ++ return self.vllm_config.kv_transfer_config is not None and self.vllm_config.kv_transfer_config.kv_connector == "DynemoNixlConnector" + @staticmethod def is_unsupported_config(engine_args: AsyncEngineArgs): @@ -2941,7 +2746,7 @@ index 85b5f31e..c501e4c8 100644 + kv_metrics.kv_active_blocks, + kv_metrics.kv_total_blocks) + -+ logger.debug("Metircs successful.") ++ logger.debug("Metircs successful.") + + except asyncio.CancelledError: + logger.debug("Shutting down MQLLMEngineClient check metrics loop.") @@ -3346,10 +3151,10 @@ index 786380c3..56a7cf89 100644 """The output data of one completion output of a request. diff --git a/vllm/remote_prefill.py b/vllm/remote_prefill.py new file mode 100644 -index 00000000..957f55de +index 00000000..03f02006 --- /dev/null +++ b/vllm/remote_prefill.py -@@ -0,0 +1,54 @@ +@@ -0,0 +1,53 @@ +from dataclasses import dataclass +from typing import Callable, Optional, List, Coroutine + @@ -3387,7 +3192,6 @@ index 00000000..957f55de + """ + request_id: str + src_block_ids: List[int] -+ staging_block_ids: List[int] + dst_block_ids: List[int] + dst_engine_id: str + notify_msg: str @@ -3531,7 +3335,7 @@ index 12baecde..cbada27f 100644 if self.vllm_config.kv_transfer_config is None: return False + -+ if self.vllm_config.kv_transfer_config.kv_connector == "TritonNixlConnector": ++ if self.vllm_config.kv_transfer_config.kv_connector == "DynemoNixlConnector": + return False prefill_meta = model_input.attn_metadata.prefill_metadata @@ -3541,13 +3345,13 @@ index 12baecde..cbada27f 100644 if self.vllm_config.kv_transfer_config is None: return False + -+ if self.vllm_config.kv_transfer_config.kv_connector == "TritonNixlConnector": ++ if self.vllm_config.kv_transfer_config.kv_connector == "DynemoNixlConnector": + return False prefill_meta = model_input.attn_metadata.prefill_metadata diff --git a/vllm/worker/worker.py b/vllm/worker/worker.py -index 582aa460..1b8515bf 100644 +index 582aa460..ffb7b403 100644 --- a/vllm/worker/worker.py +++ b/vllm/worker/worker.py @@ -2,7 +2,7 @@ @@ -3563,7 +3367,7 @@ index 582aa460..1b8515bf 100644 from vllm.worker.pooling_model_runner import PoolingModelRunner from vllm.worker.worker_base import (LocalOrDistributedWorkerBase, WorkerBase, WorkerInput) -+from vllm.distributed.device_communicators.nixl import TritonNixlConnector ++from vllm.distributed.device_communicators.nixl import DynemoNixlConnector + logger = init_logger(__name__) @@ -3577,7 +3381,7 @@ index 582aa460..1b8515bf 100644 + # TODO ptarasiewicz nixl can also support DRAM + assert self.device_config.device_type == "cuda", "Currently only CUDA is supported for Nixl connector" + -+ self.nixl_connector = TritonNixlConnector(self.vllm_config, engine_id, self.local_rank) # TODO ptarasiewicz: rank or local_rank? ++ self.nixl_connector = DynemoNixlConnector(self.vllm_config, engine_id, self.local_rank) # TODO ptarasiewicz: rank or local_rank? + assert len(self.cache_engine) == 1, "Only one cache engine is supported for now" + self.nixl_connector.register_kv_caches(self.cache_engine[0].gpu_cache) + return self.nixl_connector.agent_name @@ -3588,13 +3392,13 @@ index 582aa460..1b8515bf 100644 + + def add_remote_nixl_metadata(self, engine_id: str, agents_metadata: List[bytes], kv_caches_base_addr: List[List[Tuple[int, int]]]) -> str: + assert self.nixl_connector is not None, "Nixl connector is not initialized" -+ agent_name = self.nixl_connector.add_remote_agent(engine_id, agents_metadata, len(agents_metadata)) # TODO ptarasiewicz: rank or local_rank? -+ self.nixl_connector.add_remote_kv_caches_base_addr(engine_id, kv_caches_base_addr) ++ agent_name = self.nixl_connector.add_remote_agent(engine_id, agents_metadata[self.local_rank]) # TODO ptarasiewicz: rank or local_rank? ++ self.nixl_connector.add_remote_kv_caches_base_addr(engine_id, kv_caches_base_addr[self.local_rank]) + return agent_name + + def transfer_nixl_memory(self, src_descs: List[bytes], dst_descs: List[bytes], remote_agent_name: List[str], notify_msg: str) -> None: + assert self.nixl_connector is not None, "Nixl connector is not initialized" -+ self.nixl_connector.transfer_mem(src_descs[self.local_rank], dst_descs[self.local_rank], remote_agent_name, notify_msg) # TODO ptarasiewicz: rank or local_rank? ++ self.nixl_connector.transfer_mem(src_descs[self.local_rank], dst_descs[self.local_rank], remote_agent_name[self.local_rank], notify_msg) # TODO ptarasiewicz: rank or local_rank? + + def get_nixl_kv_caches_base_addr(self) -> List[bytes]: + assert self.nixl_connector is not None, "Nixl connector is not initialized" @@ -3602,8 +3406,8 @@ index 582aa460..1b8515bf 100644 + + def _transfer_blocks(self, worker_input: WorkerInput) -> None: + if worker_input.src_block_ids is not None: -+ for src_block_ids, staging_block_ids, dst_block_ids, dst_engine_id, notify_msg in zip(worker_input.src_block_ids, worker_input.staging_block_ids, worker_input.dst_block_ids, worker_input.dst_engine_id, worker_input.notify_msg): -+ self.nixl_connector.transfer_mem(src_block_ids, staging_block_ids, dst_block_ids, dst_engine_id, notify_msg) ++ for src_block_ids, dst_block_ids, dst_engine_id, notify_msg in zip(worker_input.src_block_ids, worker_input.dst_block_ids, worker_input.dst_engine_id, worker_input.notify_msg): ++ self.nixl_connector.transfer_mem(src_block_ids, dst_block_ids, dst_engine_id, notify_msg) + + def shutdown_nixl(self) -> None: + assert self.nixl_connector is not None, "Nixl connector is not initialized" @@ -3621,12 +3425,11 @@ index 582aa460..1b8515bf 100644 return WorkerInput( num_seq_groups=num_seq_groups, -@@ -375,6 +416,11 @@ class Worker(LocalOrDistributedWorkerBase): +@@ -375,6 +416,10 @@ class Worker(LocalOrDistributedWorkerBase): blocks_to_copy=blocks_to_copy, virtual_engine=virtual_engine, num_steps=num_steps, + src_block_ids=[r.src_block_ids for r in mem_transfer_reqs], -+ staging_block_ids=[r.staging_block_ids for r in mem_transfer_reqs], + dst_block_ids=[r.dst_block_ids for r in mem_transfer_reqs], + dst_engine_id=[r.dst_engine_id for r in mem_transfer_reqs], + notify_msg=[r.notify_msg for r in mem_transfer_reqs], @@ -3634,7 +3437,7 @@ index 582aa460..1b8515bf 100644 @torch.inference_mode() diff --git a/vllm/worker/worker_base.py b/vllm/worker/worker_base.py -index 819b81fb..ecb68530 100644 +index 819b81fb..d9c039eb 100644 --- a/vllm/worker/worker_base.py +++ b/vllm/worker/worker_base.py @@ -9,6 +9,7 @@ from typing import Any, Dict, List, Optional, Set, Tuple, Type, Union @@ -3649,7 +3452,7 @@ index 819b81fb..ecb68530 100644 from vllm.worker.model_runner_base import (BroadcastableModelInput, ModelRunnerBase, ModelRunnerInputBase) -+from vllm.distributed.device_communicators.nixl import TritonNixlConnector ++from vllm.distributed.device_communicators.nixl import DynemoNixlConnector logger = init_logger(__name__) @@ -3657,17 +3460,16 @@ index 819b81fb..ecb68530 100644 from vllm.platforms import current_platform self.current_platform = current_platform -+ self.nixl_connector: Optional[TritonNixlConnector] = None ++ self.nixl_connector: Optional[DynemoNixlConnector] = None + @abstractmethod def init_device(self) -> None: """Initialize device state, such as loading the model or other on-device -@@ -216,6 +220,12 @@ class WorkerInput: +@@ -216,6 +220,11 @@ class WorkerInput: virtual_engine: int = 0 num_steps: int = 1 + src_block_ids: Optional[List[List[int]]] = None -+ staging_block_ids: Optional[List[List[int]]] = None + dst_block_ids: Optional[List[List[int]]] = None + dst_engine_id: Optional[List[str]] = None + notify_msg: Optional[List[str]] = None @@ -3675,31 +3477,29 @@ index 819b81fb..ecb68530 100644 @classmethod def from_broadcasted_tensor_dict( cls: Type["WorkerInput"], -@@ -232,6 +242,11 @@ class WorkerInput: +@@ -232,6 +241,10 @@ class WorkerInput: blocks_to_copy=tensor_dict.pop("blocks_to_copy"), virtual_engine=tensor_dict["virtual_engine"], num_steps=tensor_dict.pop("num_steps"), + src_block_ids=tensor_dict.pop("src_block_ids"), -+ staging_block_ids=tensor_dict.pop("staging_block_ids"), + dst_block_ids=tensor_dict.pop("dst_block_ids"), + dst_engine_id=tensor_dict.pop("dst_engine_id"), + notify_msg=tensor_dict.pop("notify_msg"), ) def as_broadcastable_tensor_dict( -@@ -246,6 +261,11 @@ class WorkerInput: +@@ -246,6 +259,10 @@ class WorkerInput: "blocks_to_copy": self.blocks_to_copy, "virtual_engine": self.virtual_engine, "num_steps": self.num_steps, + "src_block_ids": self.src_block_ids, -+ "staging_block_ids": self.staging_block_ids, + "dst_block_ids": self.dst_block_ids, + "dst_engine_id": self.dst_engine_id, + "notify_msg": self.notify_msg, } return tensor_dict -@@ -316,13 +336,16 @@ class LocalOrDistributedWorkerBase(WorkerBase): +@@ -316,13 +333,16 @@ class LocalOrDistributedWorkerBase(WorkerBase): return None worker_input = WorkerInput.from_broadcasted_tensor_dict(broadcast_data) @@ -3721,7 +3521,7 @@ index 819b81fb..ecb68530 100644 def _get_driver_input_and_broadcast( self, execute_model_req: ExecuteModelRequest -@@ -396,49 +419,87 @@ class LocalOrDistributedWorkerBase(WorkerBase): +@@ -396,49 +416,79 @@ class LocalOrDistributedWorkerBase(WorkerBase): self.execute_worker(worker_input) # If there is no input, we don't need to execute the model. @@ -3818,7 +3618,7 @@ index 819b81fb..ecb68530 100644 + else: + for i in range(1, get_tp_group().world_size): + all_new_notifs.append(get_tp_group().recv_object(src=i)) - ++ + request_notif_counter = defaultdict(int) + for notifs in all_new_notifs: + for req_ids in notifs.values(): @@ -3827,20 +3627,12 @@ index 819b81fb..ecb68530 100644 + + if request_notif_counter: + logger.debug("Request notif counter: %s", request_notif_counter) -+ -+ request_done_counter = defaultdict(int) -+ for req_id in self.nixl_connector.get_done_tranfers(): -+ request_done_counter[req_id] += 1 -+ -+ if request_done_counter: -+ logger.debug("Request done counter: %s", request_done_counter) -+ + else: + request_notif_counter = {} -+ request_done_counter = {} + # output is List[SamplerOutput] - return output -+ return output, request_notif_counter, request_done_counter ++ return output, request_notif_counter + + def _transfer_blocks(self, worker_input: WorkerInput) -> None: + pass diff --git a/deploy/compoundai/sdk/src/compoundai/cli/serve_nova.py b/deploy/compoundai/sdk/src/compoundai/cli/serve_nova.py index c610fa07faa5..8ee63a810a65 100644 --- a/deploy/compoundai/sdk/src/compoundai/cli/serve_nova.py +++ b/deploy/compoundai/sdk/src/compoundai/cli/serve_nova.py @@ -26,7 +26,8 @@ from typing import Any import click -from triton_distributed_rs import DistributedRuntime, triton_endpoint, triton_worker + +from dynemo.runtime import DistributedRuntime, dynemo_endpoint, dynemo_worker logger = logging.getLogger("compoundai.serve.nova") @@ -102,7 +103,7 @@ def main( server_context.worker_index = worker_id class_instance = service.inner() - @triton_worker() + @dynemo_worker() async def worker(runtime: DistributedRuntime): if service_name and service_name != service.name: server_context.service_type = "service" @@ -157,12 +158,12 @@ async def worker(runtime: DistributedRuntime): # Bind an instance of inner to the endpoint bound_method = endpoint.func.__get__(class_instance) # Only pass request type for now, use Any for response - # TODO: Handle a triton_endpoint not having types + # TODO: Handle a dynemo_endpoint not having types # TODO: Handle multiple endpoints in a single component - triton_wrapped_method = triton_endpoint(endpoint.request_type, Any)( + dynemo_wrapped_method = dynemo_endpoint(endpoint.request_type, Any)( bound_method ) - result = await td_endpoint.serve_endpoint(triton_wrapped_method) + result = await td_endpoint.serve_endpoint(dynemo_wrapped_method) # WARNING: unreachable code :( because serve blocks logger.info(f"[{run_id}] Result: {result}") logger.info(f"[{run_id}] Registered endpoint '{name}'") diff --git a/deploy/compoundai/sdk/src/compoundai/sdk/decorators.py b/deploy/compoundai/sdk/src/compoundai/sdk/decorators.py index 43328bc75515..6ef09f661933 100644 --- a/deploy/compoundai/sdk/src/compoundai/sdk/decorators.py +++ b/deploy/compoundai/sdk/src/compoundai/sdk/decorators.py @@ -50,7 +50,7 @@ async def __call__(self, *args: t.Any, **kwargs: t.Any) -> Any: if isinstance(args[1], (str, dict)): args[1] = self.request_type.parse_obj(args[1]) # type: ignore - # Convert Pydantic model to dict before passing to triton + # Convert Pydantic model to dict before passing to dynemo if len(args) > 1 and isinstance(args[1], BaseModel): args = list(args) # type: ignore args[1] = args[1].model_dump() # type: ignore diff --git a/deploy/compoundai/sdk/src/compoundai/sdk/dependency.py b/deploy/compoundai/sdk/src/compoundai/sdk/dependency.py index b58aec6dfb49..de3d62520cc9 100644 --- a/deploy/compoundai/sdk/src/compoundai/sdk/dependency.py +++ b/deploy/compoundai/sdk/src/compoundai/sdk/dependency.py @@ -72,9 +72,9 @@ async def stream_worker(): else: # Create nova worker if no runtime - from triton_distributed_rs import DistributedRuntime, triton_worker + from dynemo.runtime import DistributedRuntime, dynemo_worker - @triton_worker() + @dynemo_worker() async def stream_worker(runtime: DistributedRuntime): try: # Store runtime for future use diff --git a/examples/python_rs/llm/tensorrt_llm/README.md b/examples/python_rs/llm/tensorrt_llm/README.md index 3a592f44eed7..ebf269536da7 100644 --- a/examples/python_rs/llm/tensorrt_llm/README.md +++ b/examples/python_rs/llm/tensorrt_llm/README.md @@ -90,14 +90,14 @@ Note: NATS and ETCD servers should be running and accessible from the container Run the server logging (with debug level logging): ```bash -TRD_LOG=DEBUG http & +DYN_LOG=DEBUG http & ``` By default the server will run on port 8080. Add model to the server: ```bash -llmctl http add chat TinyLlama/TinyLlama-1.1B-Chat-v1.0 triton-init.tensorrt-llm.chat/completions -llmctl http add completion TinyLlama/TinyLlama-1.1B-Chat-v1.0 triton-init.tensorrt-llm.completions +llmctl http add chat TinyLlama/TinyLlama-1.1B-Chat-v1.0 dynemo.tensorrt-llm.chat/completions +llmctl http add completion TinyLlama/TinyLlama-1.1B-Chat-v1.0 dynemo.tensorrt-llm.completions ``` #### 2. Workers @@ -214,14 +214,14 @@ Run the container interactively with the following command: Run the server logging (with debug level logging): ```bash -TRD_LOG=DEBUG http & +DYN_LOG=DEBUG http & ``` By default the server will run on port 8080. Add model to the server: ```bash -llmctl http add chat TinyLlama/TinyLlama-1.1B-Chat-v1.0 triton-init.router.chat/completions -llmctl http add completion TinyLlama/TinyLlama-1.1B-Chat-v1.0 triton-init.router.completions +llmctl http add chat TinyLlama/TinyLlama-1.1B-Chat-v1.0 dynemo.router.chat/completions +llmctl http add completion TinyLlama/TinyLlama-1.1B-Chat-v1.0 dynemo.router.completions ``` #### 2. Workers diff --git a/examples/python_rs/llm/tensorrt_llm/common/client.py b/examples/python_rs/llm/tensorrt_llm/common/client.py index 5ca4a53d3882..22bd43e7781a 100644 --- a/examples/python_rs/llm/tensorrt_llm/common/client.py +++ b/examples/python_rs/llm/tensorrt_llm/common/client.py @@ -19,12 +19,12 @@ import uvloop -from triton_distributed.runtime import DistributedRuntime, triton_worker +from dynemo.runtime import DistributedRuntime, dynemo_worker from .protocol import Request -@triton_worker() +@dynemo_worker() async def worker( runtime: DistributedRuntime, component: str, @@ -38,7 +38,7 @@ async def worker( """ # create client client = ( - await runtime.namespace("triton-init") + await runtime.namespace("dynemo") .component(component) .endpoint("generate") .client() diff --git a/examples/python_rs/llm/tensorrt_llm/disaggregated/router.py b/examples/python_rs/llm/tensorrt_llm/disaggregated/router.py index 5511ed04a48f..9214a8d58b0c 100644 --- a/examples/python_rs/llm/tensorrt_llm/disaggregated/router.py +++ b/examples/python_rs/llm/tensorrt_llm/disaggregated/router.py @@ -26,11 +26,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.serve.openai_protocol import CompletionRequest, DisaggregatedParams -from triton_distributed.runtime import ( - DistributedRuntime, - triton_endpoint, - triton_worker, -) +from dynemo.runtime import DistributedRuntime, dynemo_endpoint, dynemo_worker logger.set_level("debug") @@ -73,7 +69,7 @@ async def _get_ctx_resp(self, request, ctx_client): # Disagg params should be in under the choices field in the response object. # This is the case for completions but not for chat. - @triton_endpoint(CompletionRequest, DisaggCompletionStreamResponse) + @dynemo_endpoint(CompletionRequest, DisaggCompletionStreamResponse) async def generate_completion(self, request): # These settings are needed to satisfy request checks. request.skip_special_tokens = False @@ -106,7 +102,7 @@ async def generate_completion(self, request): ) yield json.loads(gen_resp_obj.model_dump_json(exclude_unset=True)) - @triton_endpoint(DisaggChatCompletionRequest, DisaggChatCompletionStreamResponse) + @dynemo_endpoint(DisaggChatCompletionRequest, DisaggChatCompletionStreamResponse) async def generate_chat(self, request): # These settings are needed to satisfy request checks. request.skip_special_tokens = False @@ -140,35 +136,35 @@ async def generate_chat(self, request): yield json.loads(gen_resp_obj.model_dump_json(exclude_unset=True)) -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime): """ Instantiate a `backend` component and serve the `generate` endpoint A `Component` can serve multiple endpoints """ - component = runtime.namespace("triton-init").component("router") + component = runtime.namespace("dynemo").component("router") await component.create_service() ctx_completion_client = ( - await runtime.namespace("triton-init") + await runtime.namespace("dynemo") .component("tensorrt-llm-ctx") .endpoint("completions") .client() ) gen_completion_client = ( - await runtime.namespace("triton-init") + await runtime.namespace("dynemo") .component("tensorrt-llm-gen") .endpoint("completions") .client() ) ctx_chat_client = ( - await runtime.namespace("triton-init") + await runtime.namespace("dynemo") .component("tensorrt-llm-ctx") .endpoint("chat/completions") .client() ) gen_chat_client = ( - await runtime.namespace("triton-init") + await runtime.namespace("dynemo") .component("tensorrt-llm-gen") .endpoint("chat/completions") .client() diff --git a/examples/python_rs/llm/tensorrt_llm/disaggregated/worker.py b/examples/python_rs/llm/tensorrt_llm/disaggregated/worker.py index 573f64f99e7f..f2f0f5c41d60 100644 --- a/examples/python_rs/llm/tensorrt_llm/disaggregated/worker.py +++ b/examples/python_rs/llm/tensorrt_llm/disaggregated/worker.py @@ -44,11 +44,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.serve.openai_protocol import CompletionRequest -from triton_distributed.runtime import ( - DistributedRuntime, - triton_endpoint, - triton_worker, -) +from dynemo.runtime import DistributedRuntime, dynemo_endpoint, dynemo_worker logger.set_level("debug") @@ -89,7 +85,7 @@ def __init__( engine_config.extra_args["_mpi_session"] = self._mpi_session super().__init__(engine_config) - @triton_endpoint(DisaggChatCompletionRequest, DisaggChatCompletionStreamResponse) + @dynemo_endpoint(DisaggChatCompletionRequest, DisaggChatCompletionStreamResponse) async def generate_chat(self, request): if self._llm_engine is None: raise RuntimeError("Engine not initialized") @@ -168,7 +164,7 @@ async def generate_chat(self, request): self._ongoing_request_count -= 1 - @triton_endpoint(CompletionRequest, DisaggCompletionStreamResponse) + @dynemo_endpoint(CompletionRequest, DisaggCompletionStreamResponse) async def generate_completions(self, request): if self._llm_engine is None: raise RuntimeError("Engine not initialized") @@ -215,7 +211,7 @@ async def generate_completions(self, request): self._ongoing_request_count -= 1 -@triton_worker() +@dynemo_worker() async def worker( runtime: DistributedRuntime, engine_config: LLMAPIConfig, @@ -230,9 +226,7 @@ async def worker( server_type = disagg_config.server_configs[instance_idx].type logger.info(f"Starting {server_type} server") - component = runtime.namespace("triton-init").component( - f"tensorrt-llm-{server_type}" - ) + component = runtime.namespace("dynemo").component(f"tensorrt-llm-{server_type}") await component.create_service() completions_endpoint = component.endpoint("completions") diff --git a/examples/python_rs/llm/tensorrt_llm/monolith/worker.py b/examples/python_rs/llm/tensorrt_llm/monolith/worker.py index 162970289502..a7d3c2f6d6a0 100644 --- a/examples/python_rs/llm/tensorrt_llm/monolith/worker.py +++ b/examples/python_rs/llm/tensorrt_llm/monolith/worker.py @@ -32,11 +32,7 @@ CompletionStreamResponse, ) -from triton_distributed.runtime import ( - DistributedRuntime, - triton_endpoint, - triton_worker, -) +from dynemo.runtime import DistributedRuntime, dynemo_endpoint, dynemo_worker logger.set_level("debug") @@ -49,7 +45,7 @@ class TensorrtLLMEngine(BaseTensorrtLLMEngine): def __init__(self, engine_config: LLMAPIConfig): super().__init__(engine_config) - @triton_endpoint(ChatCompletionRequest, ChatCompletionStreamResponse) + @dynemo_endpoint(ChatCompletionRequest, ChatCompletionStreamResponse) async def generate_chat(self, request): if self._llm_engine is None: raise RuntimeError("Engine not initialized") @@ -97,7 +93,7 @@ async def generate_chat(self, request): except Exception as e: raise RuntimeError("Failed to generate: " + str(e)) - @triton_endpoint(CompletionRequest, CompletionStreamResponse) + @dynemo_endpoint(CompletionRequest, CompletionStreamResponse) async def generate_completion(self, request): if self._llm_engine is None: raise RuntimeError("Engine not initialized") @@ -144,13 +140,13 @@ async def generate_completion(self, request): raise RuntimeError("Failed to generate: " + str(e)) -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime, engine_config: LLMAPIConfig): """ Instantiate a `backend` component and serve the `generate` endpoint A `Component` can serve multiple endpoints """ - component = runtime.namespace("triton-init").component("tensorrt-llm") + component = runtime.namespace("dynemo").component("tensorrt-llm") await component.create_service() completions_endpoint = component.endpoint("completions") diff --git a/examples/python_rs/llm/vllm/README.md b/examples/python_rs/llm/vllm/README.md index 9c5502a7e448..51b060e22313 100644 --- a/examples/python_rs/llm/vllm/README.md +++ b/examples/python_rs/llm/vllm/README.md @@ -57,13 +57,13 @@ The example is designed to run in a containerized environment using Triton Distr Run the server logging (with debug level logging): ```bash -TRD_LOG=DEBUG http +DYN_LOG=DEBUG http ``` By default the server will run on port 8080. Add model to the server: ```bash -llmctl http add chat-models deepseek-ai/DeepSeek-R1-Distill-Llama-8B triton-init.vllm.generate +llmctl http add chat-models deepseek-ai/DeepSeek-R1-Distill-Llama-8B dynemo.vllm.generate ``` ##### Example Output @@ -71,7 +71,7 @@ llmctl http add chat-models deepseek-ai/DeepSeek-R1-Distill-Llama-8B triton-init +------------+------------------------------------------+-----------+-----------+----------+ | MODEL TYPE | MODEL NAME | NAMESPACE | COMPONENT | ENDPOINT | +------------+------------------------------------------+-----------+-----------+----------+ -| chat | deepseek-ai/DeepSeek-R1-Distill-Llama-8B | triton-init | vllm | generate | +| chat | deepseek-ai/DeepSeek-R1-Distill-Llama-8B | dynemo | vllm | generate | +------------+------------------------------------------+-----------+-----------+----------+ ``` @@ -83,7 +83,7 @@ In a separate terminal run the vllm worker: ```bash # Activate virtual environment -source /opt/triton/venv/bin/activate +source /opt/dynemo/venv/bin/activate # Launch worker cd /workspace/examples/python_rs/llm/vllm @@ -116,7 +116,7 @@ This deployment option splits the model serving across prefill and decode worker **Terminal 1 - Prefill Worker:** ```bash # Activate virtual environment -source /opt/triton/venv/bin/activate +source /opt/dynemo/venv/bin/activate # Launch prefill worker cd /workspace/examples/python_rs/llm/vllm @@ -126,7 +126,7 @@ VLLM_WORKER_MULTIPROC_METHOD=spawn CUDA_VISIBLE_DEVICES=0 python3 -m disaggregat --enforce-eager \ --tensor-parallel-size 1 \ --kv-transfer-config \ - '{"kv_connector":"TritonNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2}' + '{"kv_connector":"DynemoNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2}' ``` ##### Example Output @@ -142,7 +142,7 @@ INFO 03-02 05:59:47 llm_engine.py:476] init engine (profile, create kv cache, wa **Terminal 2 - Decode Worker:** ```bash # Activate virtual environment -source /opt/triton/venv/bin/activate +source /opt/dynemo/venv/bin/activate # Launch decode worker cd /workspace/examples/python_rs/llm/vllm @@ -152,7 +152,7 @@ VLLM_WORKER_MULTIPROC_METHOD=spawn CUDA_VISIBLE_DEVICES=1,2 python3 -m disaggreg --enforce-eager \ --tensor-parallel-size 2 \ --kv-transfer-config \ - '{"kv_connector":"TritonNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2}' + '{"kv_connector":"DynemoNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2}' ``` The disaggregated deployment utilizes separate GPUs for prefill and decode operations, allowing for optimized resource allocation and improved performance. For more details on the disaggregated deployment, please refer to the [vLLM documentation](https://docs.vllm.ai/en/latest/features/disagg_prefill.html). @@ -254,7 +254,7 @@ tmux ls | grep 'v-' | cut -d: -f1 | xargs -I{} tmux kill-session -t {} **Terminal 1 - Router:** ```bash # Activate virtual environment -source /opt/triton/venv/bin/activate +source /opt/dynemo/venv/bin/activate # Launch prefill worker cd /workspace/examples/python_rs/llm/vllm @@ -270,7 +270,7 @@ You can choose only the prefix strategy for now: **Terminal 2 - Processor:** ```bash # Activate virtual environment -source /opt/triton/venv/bin/activate +source /opt/dynemo/venv/bin/activate # Processor must take the same args as the worker # This is temporary until we communicate the ModelDeploymentCard over etcd @@ -286,7 +286,7 @@ RUST_LOG=info python3 -m kv_router.processor \ **Terminal 3 and 4 - Workers:** ```bash # Activate virtual environment -source /opt/triton/venv/bin/activate +source /opt/dynemo/venv/bin/activate # Launch Worker 1 and Worker 2 with same command cd /workspace/examples/python_rs/llm/vllm @@ -304,7 +304,7 @@ Note: block-size must be 64, otherwise Router won't work (accepts only 64 tokens **Terminal 5 - Client:** Don't forget to add the model to the server: ```bash -llmctl http add chat-models deepseek-ai/DeepSeek-R1-Distill-Llama-8B triton-init.process.chat/completions +llmctl http add chat-models deepseek-ai/DeepSeek-R1-Distill-Llama-8B dynemo.process.chat/completions ``` ```bash @@ -351,7 +351,7 @@ Run following commands in 4 terminals: **Terminal 1 - vLLM Worker:** ```bash # Activate virtual environment -source /opt/triton/venv/bin/activate +source /opt/dynemo/venv/bin/activate cd /workspace/examples/python_rs/llm/vllm RUST_LOG=info python3 -m preprocessor.worker --model deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B @@ -361,7 +361,7 @@ RUST_LOG=info python3 -m preprocessor.worker --model deepseek-ai/DeepSeek-R1-Dis ```bash # Activate virtual environment -source /opt/triton/venv/bin/activate +source /opt/dynemo/venv/bin/activate cd /workspace/examples/python_rs/llm/vllm RUST_LOG=info python3 -m preprocessor.processor --model deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B @@ -371,13 +371,13 @@ RUST_LOG=info python3 -m preprocessor.processor --model deepseek-ai/DeepSeek-R1- Run the server logging (with debug level logging): ```bash -TRD_LOG=DEBUG http +DYN_LOG=DEBUG http ``` By default the server will run on port 8080. Add model to the server: ```bash -llmctl http add chat-models deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B triton-init.preprocessor.generate +llmctl http add chat-models deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B dynemo.preprocessor.generate ``` **Terminal 4 - client** diff --git a/examples/python_rs/llm/vllm/common/client.py b/examples/python_rs/llm/vllm/common/client.py index e67874c633d3..d2fb5f4368d2 100644 --- a/examples/python_rs/llm/vllm/common/client.py +++ b/examples/python_rs/llm/vllm/common/client.py @@ -19,12 +19,12 @@ import uvloop -from triton_distributed.runtime import DistributedRuntime, triton_worker +from dynemo.runtime import DistributedRuntime, dynemo_worker from .protocol import Request -@triton_worker() +@dynemo_worker() async def worker( runtime: DistributedRuntime, component: str, @@ -36,9 +36,7 @@ async def worker( Instantiate a `backend` client and call the `generate` endpoint """ # get endpoint - endpoint = ( - runtime.namespace("triton-init").component(component).endpoint("generate") - ) + endpoint = runtime.namespace("dynemo").component(component).endpoint("generate") # create client client = await endpoint.client() diff --git a/examples/python_rs/llm/vllm/disaggregated/decode_worker.py b/examples/python_rs/llm/vllm/disaggregated/decode_worker.py index a6eaf9ad3c8b..481bfe853978 100644 --- a/examples/python_rs/llm/vllm/disaggregated/decode_worker.py +++ b/examples/python_rs/llm/vllm/disaggregated/decode_worker.py @@ -31,11 +31,7 @@ ) from vllm.logger import logger as vllm_logger -from triton_distributed.runtime import ( - DistributedRuntime, - triton_endpoint, - triton_worker, -) +from dynemo.runtime import DistributedRuntime, dynemo_endpoint, dynemo_worker class VllmDecodeEngine(BaseVllmEngine, ProcessMixIn): @@ -58,7 +54,7 @@ def __init__(self, engine_args: AsyncEngineArgs, prefill): self.kv_transfer_config = engine_args.create_engine_config().kv_transfer_config self.kv_rank = self.kv_transfer_config.kv_rank - @triton_endpoint(ChatCompletionRequest, ChatCompletionStreamResponse) + @dynemo_endpoint(ChatCompletionRequest, ChatCompletionStreamResponse) async def generate(self, raw_request): if self.engine_client is None: await self.initialize() @@ -107,17 +103,17 @@ async def generate(self, raw_request): await prefill_output -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime, engine_args: AsyncEngineArgs): """ Instantiate a `backend` component and serve the `generate` endpoint A `Component` can serve multiple endpoints """ - component = runtime.namespace("triton-init").component("vllm") + component = runtime.namespace("dynemo").component("vllm") await component.create_service() prefill = ( - await runtime.namespace("triton-init") + await runtime.namespace("dynemo") .component("prefill") .endpoint("generate") .client() diff --git a/examples/python_rs/llm/vllm/disaggregated/prefill_worker.py b/examples/python_rs/llm/vllm/disaggregated/prefill_worker.py index 58e399f58798..423550af40ac 100644 --- a/examples/python_rs/llm/vllm/disaggregated/prefill_worker.py +++ b/examples/python_rs/llm/vllm/disaggregated/prefill_worker.py @@ -24,11 +24,7 @@ from vllm.engine.arg_utils import AsyncEngineArgs from vllm.logger import logger as vllm_logger -from triton_distributed.runtime import ( - DistributedRuntime, - triton_endpoint, - triton_worker, -) +from dynemo.runtime import DistributedRuntime, dynemo_endpoint, dynemo_worker class VllmPrefillEngine(BaseVllmEngine): @@ -49,7 +45,7 @@ def __init__(self, engine_args: AsyncEngineArgs): self.kv_transfer_config = engine_args.create_engine_config().kv_transfer_config self.kv_rank = self.kv_transfer_config.kv_rank - @triton_endpoint(PrefillRequest, PrefillResponse) + @dynemo_endpoint(PrefillRequest, PrefillResponse) async def generate(self, request): if self.engine_client is None: await self.initialize() @@ -66,13 +62,13 @@ async def generate(self, request): yield True -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime, engine_args: AsyncEngineArgs): """ Instantiate a `backend` component and serve the `generate` endpoint A `Component` can serve multiple endpoints """ - component = runtime.namespace("triton-init").component("prefill") + component = runtime.namespace("dynemo").component("prefill") await component.create_service() async with VllmPrefillEngine(engine_args) as prefill_engine: diff --git a/examples/python_rs/llm/vllm/kv_router/metrics_router.py b/examples/python_rs/llm/vllm/kv_router/metrics_router.py index 9913e34ad0a4..3fe27dc47ab4 100644 --- a/examples/python_rs/llm/vllm/kv_router/metrics_router.py +++ b/examples/python_rs/llm/vllm/kv_router/metrics_router.py @@ -20,12 +20,8 @@ from common.protocol import Request, Response from vllm.logger import logger as vllm_logger -from triton_distributed.llm import KvRouter -from triton_distributed.runtime import ( - DistributedRuntime, - triton_endpoint, - triton_worker, -) +from dynemo.llm import KvRouter +from dynemo.runtime import DistributedRuntime, dynemo_endpoint, dynemo_worker class Router: @@ -41,7 +37,7 @@ def __init__( self.router = router self.workers_client = workers_client - @triton_endpoint(Request, Response) + @dynemo_endpoint(Request, Response) async def generate(self, request): lora_id = 0 worker_id = None @@ -73,7 +69,7 @@ async def generate(self, request): resp = resp.data() if hasattr(resp, "data") else resp yield resp - @triton_endpoint(Request, Response) + @dynemo_endpoint(Request, Response) async def mock_generate(self, request): print(f"Received request: {request}") yield "Hello, World!" @@ -82,10 +78,10 @@ async def mock_generate(self, request): ROUTE_SELF = True -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime): workers_client = ( - await runtime.namespace("triton-init") + await runtime.namespace("dynemo") .component("vllm") .endpoint("generate") .client() @@ -102,7 +98,7 @@ async def worker(runtime: DistributedRuntime): # simply be ignored, but before that, we will make sure that the services # of the same namespace::component are created via KvMetricsPublisher, # if it is also used to create endpoints. - kv_listener = runtime.namespace("triton-init").component("vllm") + kv_listener = runtime.namespace("dynemo").component("vllm") await kv_listener.create_service() router = KvRouter(runtime, kv_listener) # i.e. below will cause panic @@ -111,7 +107,7 @@ async def worker(runtime: DistributedRuntime): # Router(router, workers_client).mock_generate # ) - router_component = runtime.namespace("triton-init").component("frontend") + router_component = runtime.namespace("dynemo").component("frontend") await router_component.create_service() endpoint = router_component.endpoint("generate") diff --git a/examples/python_rs/llm/vllm/kv_router/mock_worker.py b/examples/python_rs/llm/vllm/kv_router/mock_worker.py index 75980443b065..26e1bd808da2 100644 --- a/examples/python_rs/llm/vllm/kv_router/mock_worker.py +++ b/examples/python_rs/llm/vllm/kv_router/mock_worker.py @@ -22,15 +22,11 @@ from common.protocol import Request, Response from vllm.logger import logger as vllm_logger -from triton_distributed.llm import KvMetricsPublisher -from triton_distributed.runtime import ( - DistributedRuntime, - triton_endpoint, - triton_worker, -) +from dynemo.llm import KvMetricsPublisher +from dynemo.runtime import DistributedRuntime, dynemo_endpoint, dynemo_worker -class TritonResult: +class DynemoResult: OK = 0 ERR = 1 @@ -43,19 +39,17 @@ class MockEngine: def __init__(self, metrics_publisher, worker_id): self.worker_id = worker_id # KV events - self.lib = ctypes.CDLL("/opt/triton/llm_binding/lib/libtriton_llm_capi.so") - self.lib.triton_llm_init.argtypes = [c_char_p, c_char_p, c_int64] - self.lib.triton_llm_init.restype = c_uint32 - result = self.lib.triton_llm_init( - "triton-init".encode(), "vllm".encode(), worker_id - ) - if result == TritonResult.OK: + self.lib = ctypes.CDLL("/opt/dynemo/llm_binding/lib/libdynemo_llm_capi.so") + self.lib.dynemo_llm_init.argtypes = [c_char_p, c_char_p, c_int64] + self.lib.dynemo_llm_init.restype = c_uint32 + result = self.lib.dynemo_llm_init("dynemo".encode(), "vllm".encode(), worker_id) + if result == DynemoResult.OK: vllm_logger.info( "KVCacheEventManager initialized successfully. Ready to publish KV Cache Events" ) else: vllm_logger.info("KVCacheEventManager initialization failed!") - self.lib.triton_kv_event_publish_stored.argtypes = [ + self.lib.dynemo_kv_event_publish_stored.argtypes = [ ctypes.c_uint64, # event_id ctypes.POINTER(ctypes.c_uint32), # token_ids ctypes.POINTER(ctypes.c_size_t), # num_block_tokens @@ -64,18 +58,18 @@ def __init__(self, metrics_publisher, worker_id): ctypes.POINTER(ctypes.c_uint64), # parent_hash ctypes.c_uint64, # lora_id ] - self.lib.triton_kv_event_publish_stored.restype = ( + self.lib.dynemo_kv_event_publish_stored.restype = ( ctypes.c_uint32 - ) # triton_llm_result_t + ) # dynemo_llm_result_t - self.lib.triton_kv_event_publish_removed.argtypes = [ + self.lib.dynemo_kv_event_publish_removed.argtypes = [ ctypes.c_uint64, # event_id ctypes.POINTER(ctypes.c_uint64), # block_ids ctypes.c_size_t, # num_blocks ] - self.lib.triton_kv_event_publish_removed.restype = ( + self.lib.dynemo_kv_event_publish_removed.restype = ( ctypes.c_uint32 - ) # triton_llm_result_t + ) # dynemo_llm_result_t # KV metrics self.metrics_publisher = metrics_publisher @@ -95,7 +89,7 @@ def __init__(self, metrics_publisher, worker_id): self.event_id_counter = 0 self.tokens = [3] * 64 - @triton_endpoint(Request, Response) + @dynemo_endpoint(Request, Response) async def generate(self, request): print(f"Received request: {request}") self.request_active_slots = min( @@ -117,7 +111,7 @@ def store_event(self): if self.event_id_counter > 0 else None ) - result = self.lib.triton_kv_event_publish_stored( + result = self.lib.dynemo_kv_event_publish_stored( self.event_id_counter, # uint64_t event_id (ctypes.c_uint32 * len(self.tokens))( *self.tokens @@ -132,7 +126,7 @@ def store_event(self): ) self.event_id_counter += 1 - if result == TritonResult.OK: + if result == DynemoResult.OK: vllm_logger.debug(f"Store - Published KV Event: {self.event_id_counter}") else: vllm_logger.debug( @@ -152,13 +146,13 @@ async def cooldown(self): ) -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime): """ Instantiate a `backend` component and serve the `generate` endpoint A `Component` can serve multiple endpoints """ - component = runtime.namespace("triton-init").component("vllm") + component = runtime.namespace("dynemo").component("vllm") metrics_publisher = KvMetricsPublisher() await metrics_publisher.create_service(component) diff --git a/examples/python_rs/llm/vllm/kv_router/processor.py b/examples/python_rs/llm/vllm/kv_router/processor.py index bcb7cdeaf92c..e10e130786ef 100644 --- a/examples/python_rs/llm/vllm/kv_router/processor.py +++ b/examples/python_rs/llm/vllm/kv_router/processor.py @@ -34,12 +34,7 @@ from vllm.outputs import RequestOutput from vllm.transformers_utils.tokenizer import AnyTokenizer -from triton_distributed.runtime import ( - Client, - DistributedRuntime, - triton_endpoint, - triton_worker, -) +from dynemo.runtime import Client, DistributedRuntime, dynemo_endpoint, dynemo_worker class RequestType(Enum): @@ -162,38 +157,38 @@ async def _generate_responses( f"Request type {request_type} not implemented" ) - @triton_endpoint(ChatCompletionRequest, ChatCompletionStreamResponse) + @dynemo_endpoint(ChatCompletionRequest, ChatCompletionStreamResponse) async def generate_chat(self, raw_request): async for response in self._generate(raw_request, RequestType.CHAT): yield response - @triton_endpoint(CompletionRequest, CompletionStreamResponse) + @dynemo_endpoint(CompletionRequest, CompletionStreamResponse) async def generate_completions(self, raw_request): async for response in self._generate(raw_request, RequestType.COMPLETION): yield response -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime, engine_args: AsyncEngineArgs): """ Set up clients to the router and workers. - Serve the triton-init.process.chat/completions endpoint. + Serve the dynemo.process.chat/completions endpoint. """ workers_client = ( - await runtime.namespace("triton-init") + await runtime.namespace("dynemo") .component("vllm") .endpoint("generate") .client() ) router_client = ( - await runtime.namespace("triton-init") + await runtime.namespace("dynemo") .component("router") .endpoint("generate") .client() ) - preprocess_component = runtime.namespace("triton-init").component("process") + preprocess_component = runtime.namespace("dynemo").component("process") await preprocess_component.create_service() chat_endpoint = preprocess_component.endpoint("chat/completions") diff --git a/examples/python_rs/llm/vllm/kv_router/router.py b/examples/python_rs/llm/vllm/kv_router/router.py index db3f41e4ad23..29b9b0636444 100644 --- a/examples/python_rs/llm/vllm/kv_router/router.py +++ b/examples/python_rs/llm/vllm/kv_router/router.py @@ -23,12 +23,8 @@ from common.protocol import Tokens from vllm.logger import logger as vllm_logger -from triton_distributed.llm import KvRouter -from triton_distributed.runtime import ( - DistributedRuntime, - triton_endpoint, - triton_worker, -) +from dynemo.llm import KvRouter +from dynemo.runtime import DistributedRuntime, dynemo_endpoint, dynemo_worker WorkerId = str @@ -55,7 +51,7 @@ def __init__( self.router = router self.routing_strategy = routing_strategy - @triton_endpoint(Tokens, WorkerId) + @dynemo_endpoint(Tokens, WorkerId) async def generate(self, request) -> AsyncIterator[WorkerId]: lora_id = 0 worker_id = None @@ -82,14 +78,14 @@ async def generate(self, request) -> AsyncIterator[WorkerId]: ) -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime, args: Namespace): """ Set up the worker clients. - Serve the triton-init.router.generate endpoint. + Serve the dynemo.router.generate endpoint. """ workers_client = ( - await runtime.namespace("triton-init") + await runtime.namespace("dynemo") .component("vllm") .endpoint("generate") .client() @@ -114,10 +110,10 @@ async def worker(runtime: DistributedRuntime, args: Namespace): + "\n".join(f"id: {id}" for id in workers_client.endpoint_ids()) ) - kv_listener = runtime.namespace("triton-init").component("vllm") + kv_listener = runtime.namespace("dynemo").component("vllm") await kv_listener.create_service() - router_component = runtime.namespace("triton-init").component("router") + router_component = runtime.namespace("dynemo").component("router") await router_component.create_service() router = KvRouter(runtime, kv_listener) diff --git a/examples/python_rs/llm/vllm/kv_router/worker.py b/examples/python_rs/llm/vllm/kv_router/worker.py index 85f0933d80ad..ff44c0e98f30 100644 --- a/examples/python_rs/llm/vllm/kv_router/worker.py +++ b/examples/python_rs/llm/vllm/kv_router/worker.py @@ -25,12 +25,8 @@ from vllm.logger import logger as vllm_logger from vllm.sampling_params import RequestOutputKind -from triton_distributed.llm import KvMetricsPublisher -from triton_distributed.runtime import ( - DistributedRuntime, - triton_endpoint, - triton_worker, -) +from dynemo.llm import KvMetricsPublisher +from dynemo.runtime import DistributedRuntime, dynemo_endpoint, dynemo_worker vllm_logger.info(f"VLLM_KV_CAPI_PATH: {os.environ['VLLM_KV_CAPI_PATH']}") @@ -52,7 +48,7 @@ async def initialize(self): assert self.engine_client is not None, "engine_client was not initialized" self.engine_client.set_metrics_publisher(self.metrics_publisher) - @triton_endpoint(vLLMGenerateRequest, MyRequestOutput) + @dynemo_endpoint(vLLMGenerateRequest, MyRequestOutput) async def generate(self, request) -> AsyncIterator: assert ( self.engine_client is not None @@ -77,12 +73,12 @@ async def generate(self, request) -> AsyncIterator: ).model_dump_json() -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime, engine_args: AsyncEngineArgs): """ - Serve the triton-init.vllm.generate endpoint. + Serve the dynemo.vllm.generate endpoint. """ - worker_component = runtime.namespace("triton-init").component("vllm") + worker_component = runtime.namespace("dynemo").component("vllm") await worker_component.create_service() worker_endpoint = worker_component.endpoint("generate") @@ -91,7 +87,7 @@ async def worker(runtime: DistributedRuntime, engine_args: AsyncEngineArgs): os.environ["VLLM_WORKER_ID"] = str(VLLM_WORKER_ID) vllm_logger.info(f"Generate endpoint ID: {VLLM_WORKER_ID}") - VLLM_KV_NAMESPACE = "triton-init" + VLLM_KV_NAMESPACE = "dynemo" os.environ["VLLM_KV_NAMESPACE"] = str(VLLM_KV_NAMESPACE) VLLM_KV_COMPONENT = "vllm" diff --git a/examples/python_rs/llm/vllm/monolith/worker.py b/examples/python_rs/llm/vllm/monolith/worker.py index 5ec395b7b848..28cfeee1c76b 100644 --- a/examples/python_rs/llm/vllm/monolith/worker.py +++ b/examples/python_rs/llm/vllm/monolith/worker.py @@ -28,11 +28,7 @@ ) from vllm.logger import logger as vllm_logger -from triton_distributed.runtime import ( - DistributedRuntime, - triton_endpoint, - triton_worker, -) +from dynemo.runtime import DistributedRuntime, dynemo_endpoint, dynemo_worker class VllmEngine(BaseVllmEngine, ProcessMixIn): @@ -43,7 +39,7 @@ class VllmEngine(BaseVllmEngine, ProcessMixIn): def __init__(self, engine_args: AsyncEngineArgs): super().__init__(engine_args) - @triton_endpoint(ChatCompletionRequest, ChatCompletionStreamResponse) + @dynemo_endpoint(ChatCompletionRequest, ChatCompletionStreamResponse) async def generate(self, raw_request): if self.engine_client is None: await self.initialize() @@ -75,13 +71,13 @@ async def generate(self, raw_request): yield response -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime, engine_args: AsyncEngineArgs): """ Instantiate a `backend` component and serve the `generate` endpoint A `Component` can serve multiple endpoints """ - component = runtime.namespace("triton-init").component("vllm") + component = runtime.namespace("dynemo").component("vllm") await component.create_service() endpoint = component.endpoint("generate") diff --git a/examples/python_rs/llm/vllm/preprocessor/processor.py b/examples/python_rs/llm/vllm/preprocessor/processor.py index 23a0a1980543..f4d19c35532a 100644 --- a/examples/python_rs/llm/vllm/preprocessor/processor.py +++ b/examples/python_rs/llm/vllm/preprocessor/processor.py @@ -19,27 +19,27 @@ import uvloop from preprocessor.common import parse_vllm_args -from triton_distributed.runtime import ( +from dynemo.runtime import ( DistributedRuntime, ModelDeploymentCard, OAIChatPreprocessor, - triton_worker, + dynemo_worker, ) uvloop.install() -@triton_worker() +@dynemo_worker() async def preprocessor(runtime: DistributedRuntime, model_name: str, model_path: str): # create model deployment card mdc = await ModelDeploymentCard.from_local_path(model_path, model_name) # create preprocessor endpoint - component = runtime.namespace("triton-init").component("preprocessor") + component = runtime.namespace("dynemo").component("preprocessor") await component.create_service() endpoint = component.endpoint("generate") # create backend endpoint - backend = runtime.namespace("triton-init").component("backend").endpoint("generate") + backend = runtime.namespace("dynemo").component("backend").endpoint("generate") # start preprocessor service with next backend chat = OAIChatPreprocessor(mdc, endpoint, next=backend) diff --git a/examples/python_rs/llm/vllm/preprocessor/worker.py b/examples/python_rs/llm/vllm/preprocessor/worker.py index 44fab531d343..9f6d4c6d20fd 100644 --- a/examples/python_rs/llm/vllm/preprocessor/worker.py +++ b/examples/python_rs/llm/vllm/preprocessor/worker.py @@ -28,12 +28,12 @@ ) from vllm.outputs import CompletionOutput -from triton_distributed.runtime import ( +from dynemo.runtime import ( Backend, DistributedRuntime, ModelDeploymentCard, - triton_endpoint, - triton_worker, + dynemo_endpoint, + dynemo_worker, ) finish_reason_map = { @@ -107,7 +107,7 @@ def to_sampling_params(self, request) -> SamplingParams: } return SamplingParams(**sampling_params) - @triton_endpoint(Any, CompletionOutput) + @dynemo_endpoint(Any, CompletionOutput) async def generate(self, request): state = DeltaState() request_id = str(uuid.uuid4()) @@ -122,13 +122,13 @@ async def generate(self, request): yield self.to_backend_output(choice, delta_token_ids) -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime, engine_args: NvAsyncEngineArgs): """ Instantiate a `backend` component and serve the `generate` endpoint A `Component` can serve multiple endpoints """ - component = runtime.namespace("triton-init").component("backend") + component = runtime.namespace("dynemo").component("backend") await component.create_service() endpoint = component.endpoint("generate") diff --git a/examples/python_rs/llm/vllm/scripts/disaggregated.sh b/examples/python_rs/llm/vllm/scripts/disaggregated.sh index d813997f140e..96d05d60481b 100755 --- a/examples/python_rs/llm/vllm/scripts/disaggregated.sh +++ b/examples/python_rs/llm/vllm/scripts/disaggregated.sh @@ -21,12 +21,12 @@ if [ $# -gt 2 ]; then echo "Usage: $0 [model_name] [endpoint_name]" echo "Optional: model_name (default: deepseek-ai/DeepSeek-R1-Distill-Llama-8B)" - echo "Optional: endpoint_name (default: triton-init.vllm.generate)" + echo "Optional: endpoint_name (default: dynemo.vllm.generate)" exit 1 fi MODEL_NAME=${1:-"deepseek-ai/DeepSeek-R1-Distill-Llama-8B"} -ENDPOINT_NAME=${2:-"triton-init.vllm.generate"} +ENDPOINT_NAME=${2:-"dynemo.vllm.generate"} SESSION_NAME="vllm_disagg" WORKDIR="$(dirname $0)/.." INIT_CMD="cd $WORKDIR" @@ -49,7 +49,7 @@ tmux split-window -v ######################################################## HTTP_HOST="localhost" HTTP_PORT=8080 -HTTP_CMD="TRD_LOG=DEBUG http --host ${HTTP_HOST} --port ${HTTP_PORT}" +HTTP_CMD="DYN_LOG=DEBUG http --host ${HTTP_HOST} --port ${HTTP_PORT}" tmux select-pane -t 0 tmux send-keys "$INIT_CMD && $HTTP_CMD" C-m diff --git a/examples/python_rs/llm/vllm/scripts/kv-router-run.sh b/examples/python_rs/llm/vllm/scripts/kv-router-run.sh index f2889b2e4d7c..b614f23a02eb 100755 --- a/examples/python_rs/llm/vllm/scripts/kv-router-run.sh +++ b/examples/python_rs/llm/vllm/scripts/kv-router-run.sh @@ -22,18 +22,18 @@ if [ $# -lt 2 ]; then echo "Usage: $0 [model_name] [endpoint_name]" echo "Error: Must specify at least number of workers and routing strategy" echo "Optional: model_name (default: deepseek-ai/DeepSeek-R1-Distill-Llama-8B)" - echo "Optional: endpoint_name (default: triton-init.process.chat/completions)" + echo "Optional: endpoint_name (default: dynemo.process.chat/completions)" exit 1 fi NUM_WORKERS=$1 ROUTING_STRATEGY=$2 MODEL_NAME=${3:-"deepseek-ai/DeepSeek-R1-Distill-Llama-8B"} -ENDPOINT_NAME=${4:-"triton-init.process.chat/completions"} +ENDPOINT_NAME=${4:-"dynemo.process.chat/completions"} VALID_STRATEGIES=("prefix") SESSION_NAME="v" WORKDIR="/workspace/examples/python_rs/llm/vllm" -INIT_CMD="source /opt/triton/venv/bin/activate && cd $WORKDIR" +INIT_CMD="source /opt/dynemo/venv/bin/activate && cd $WORKDIR" if [[ ! " ${VALID_STRATEGIES[@]} " =~ " ${ROUTING_STRATEGY} " ]]; then echo "Error: Invalid routing strategy. Must be one of: ${VALID_STRATEGIES[*]}" @@ -42,7 +42,7 @@ fi ######################################################## # HTTP Server ######################################################## -HTTP_CMD="TRD_LOG=DEBUG http" +HTTP_CMD="DYN_LOG=DEBUG http" tmux new-session -d -s "$SESSION_NAME-http" tmux send-keys -t "$SESSION_NAME-http" "$INIT_CMD && $HTTP_CMD" C-m diff --git a/examples/python_rs/llm/vllm_nixl/README.md b/examples/python_rs/llm/vllm_nixl/README.md index 6329596df6c4..5961a9536396 100644 --- a/examples/python_rs/llm/vllm_nixl/README.md +++ b/examples/python_rs/llm/vllm_nixl/README.md @@ -33,7 +33,7 @@ All of the commands below are run inside the same container. ## Run deployment -Add model to triton and start http server. +Add model to dynemo and start http server. In terminal 0: ``` @@ -65,7 +65,7 @@ CUDA_VISIBLE_DEVICES=0 python prefill_worker.py \ --model deepseek-ai/DeepSeek-R1-Distill-Llama-8B \ --enforce-eager \ --kv-transfer-config \ - '{"kv_connector":"TritonNixlConnector"}' + '{"kv_connector":"DynemoNixlConnector"}' ``` In terminal 2: @@ -78,7 +78,7 @@ CUDA_VISIBLE_DEVICES=1,2 python3 worker.py \ --enforce-eager \ --tensor-parallel-size 2 \ --kv-transfer-config \ - '{"kv_connector":"TritonNixlConnector"}' + '{"kv_connector":"DynemoNixlConnector"}' ``` @@ -157,7 +157,7 @@ rm -r /tmp/nixl - [x] Zero copy - [x] Conditional remote prefill - [x] Manual example with tp > 1 -- [x] Run on triton distributed runtime +- [x] Run on dynemo distributed runtime - [x] add oai http endpoint - [x] Sample only on decode, do note return remote prefill response - [x] Check if all transfers finished before moving to decode diff --git a/examples/python_rs/llm/vllm_nixl/common.py b/examples/python_rs/llm/vllm_nixl/common.py index cc3ccb6e8759..30d5e50e1e85 100644 --- a/examples/python_rs/llm/vllm_nixl/common.py +++ b/examples/python_rs/llm/vllm_nixl/common.py @@ -22,7 +22,7 @@ from vllm.engine.arg_utils import AsyncEngineArgs from vllm.utils import FlexibleArgumentParser -from triton_distributed.runtime import DistributedRuntime +from dynemo.runtime import DistributedRuntime METADATA_DIR = "/tmp/nixl" diff --git a/examples/python_rs/llm/vllm_nixl/prefill_worker.py b/examples/python_rs/llm/vllm_nixl/prefill_worker.py index 6718a3fea328..c31b40dea876 100644 --- a/examples/python_rs/llm/vllm_nixl/prefill_worker.py +++ b/examples/python_rs/llm/vllm_nixl/prefill_worker.py @@ -26,7 +26,7 @@ from vllm.inputs.data import TokensPrompt from vllm.remote_prefill import RemotePrefillParams, RemotePrefillRequest -from triton_distributed.runtime import DistributedRuntime, triton_worker +from dynemo.runtime import DistributedRuntime, dynemo_worker class RequestHandler: @@ -71,7 +71,7 @@ async def generate(self, raw_request: str): yield -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime, engine_args: AsyncEngineArgs): component = runtime.namespace("test-nixl").component("prefill") await component.create_service() diff --git a/examples/python_rs/llm/vllm_nixl/worker.py b/examples/python_rs/llm/vllm_nixl/worker.py index 0829cff7d1c3..3e27ed39e0cb 100644 --- a/examples/python_rs/llm/vllm_nixl/worker.py +++ b/examples/python_rs/llm/vllm_nixl/worker.py @@ -33,11 +33,7 @@ from vllm.entrypoints.openai.serving_models import BaseModelPath, OpenAIServingModels from vllm.remote_prefill import RemotePrefillParams, RemotePrefillRequest -from triton_distributed.runtime import ( - DistributedRuntime, - triton_endpoint, - triton_worker, -) +from dynemo.runtime import DistributedRuntime, dynemo_endpoint, dynemo_worker class RequestHandler: @@ -87,7 +83,7 @@ async def callback(request: RemotePrefillRequest): return callback - @triton_endpoint(ChatCompletionRequest, ChatCompletionStreamResponse) + @dynemo_endpoint(ChatCompletionRequest, ChatCompletionStreamResponse) async def generate(self, request): if not self.initialized: await self.init() @@ -113,7 +109,7 @@ async def generate(self, request): yield response -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime, engine_args: AsyncEngineArgs): component = runtime.namespace("test-nixl").component("vllm") await component.create_service() diff --git a/examples/rust/Cargo.lock b/examples/rust/Cargo.lock index 8ca0de4b0cbc..7db2f98f1e00 100644 --- a/examples/rust/Cargo.lock +++ b/examples/rust/Cargo.lock @@ -955,6 +955,99 @@ dependencies = [ "syn 2.0.98", ] +[[package]] +name = "dynemo-llm" +version = "0.2.1" +dependencies = [ + "anyhow", + "async-openai", + "async-stream", + "async-trait", + "axum 0.8.1", + "bindgen", + "blake3", + "bs62", + "bytes", + "chrono", + "cmake", + "derive_builder", + "dynemo-runtime", + "either", + "erased-serde", + "futures", + "galil-seiferas", + "indexmap 2.7.1", + "itertools 0.14.0", + "libc", + "minijinja", + "minijinja-contrib", + "prometheus", + "pyo3", + "regex", + "semver", + "serde", + "serde-pickle", + "serde_json", + "serde_repr", + "strum", + "thiserror 2.0.11", + "tokenizers", + "tokio", + "tokio-stream", + "tokio-util", + "toktrie", + "toktrie_hf_tokenizers", + "tracing", + "unicode-segmentation", + "uuid", + "validator", + "xxhash-rust", +] + +[[package]] +name = "dynemo-runtime" +version = "0.2.1" +dependencies = [ + "anyhow", + "async-nats", + "async-once-cell", + "async-stream", + "async-trait", + "async_zmq", + "blake3", + "bytes", + "chrono", + "derive-getters", + "derive_builder", + "educe", + "either", + "etcd-client", + "figment", + "futures", + "humantime", + "local-ip-address", + "log", + "nid", + "nix", + "nuid", + "once_cell", + "prometheus", + "rand", + "regex", + "serde", + "serde_json", + "socket2", + "thiserror 1.0.69", + "tokio", + "tokio-stream", + "tokio-util", + "tracing", + "tracing-subscriber", + "uuid", + "validator", + "xxhash-rust", +] + [[package]] name = "ed25519" version = "2.2.3" @@ -1370,7 +1463,7 @@ checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" name = "hello_world" version = "0.2.0" dependencies = [ - "triton-distributed-runtime", + "dynemo-runtime", ] [[package]] @@ -1395,11 +1488,11 @@ name = "http" version = "0.2.0" dependencies = [ "clap", + "dynemo-llm", + "dynemo-runtime", "serde", "serde_json", "tokio", - "triton-distributed-llm", - "triton-distributed-runtime", ] [[package]] @@ -1895,13 +1988,13 @@ name = "llmctl" version = "0.2.0" dependencies = [ "clap", + "dynemo-llm", + "dynemo-runtime", "serde", "serde_json", "tabled", "tokio", "tracing", - "triton-distributed-llm", - "triton-distributed-runtime", ] [[package]] @@ -3297,11 +3390,11 @@ dependencies = [ name = "service_metrics" version = "0.2.0" dependencies = [ + "dynemo-runtime", "futures", "serde", "serde_json", "tokio", - "triton-distributed-runtime", ] [[package]] @@ -4000,99 +4093,6 @@ dependencies = [ "tracing-serde", ] -[[package]] -name = "triton-distributed-llm" -version = "0.2.1" -dependencies = [ - "anyhow", - "async-openai", - "async-stream", - "async-trait", - "axum 0.8.1", - "bindgen", - "blake3", - "bs62", - "bytes", - "chrono", - "cmake", - "derive_builder", - "either", - "erased-serde", - "futures", - "galil-seiferas", - "indexmap 2.7.1", - "itertools 0.14.0", - "libc", - "minijinja", - "minijinja-contrib", - "prometheus", - "pyo3", - "regex", - "semver", - "serde", - "serde-pickle", - "serde_json", - "serde_repr", - "strum", - "thiserror 2.0.11", - "tokenizers", - "tokio", - "tokio-stream", - "tokio-util", - "toktrie", - "toktrie_hf_tokenizers", - "tracing", - "triton-distributed-runtime", - "unicode-segmentation", - "uuid", - "validator", - "xxhash-rust", -] - -[[package]] -name = "triton-distributed-runtime" -version = "0.2.1" -dependencies = [ - "anyhow", - "async-nats", - "async-once-cell", - "async-stream", - "async-trait", - "async_zmq", - "blake3", - "bytes", - "chrono", - "derive-getters", - "derive_builder", - "educe", - "either", - "etcd-client", - "figment", - "futures", - "humantime", - "local-ip-address", - "log", - "nid", - "nix", - "nuid", - "once_cell", - "prometheus", - "rand", - "regex", - "serde", - "serde_json", - "socket2", - "thiserror 1.0.69", - "tokio", - "tokio-stream", - "tokio-util", - "tracing", - "tracing-subscriber", - "uuid", - "validator", - "xxhash-rust", -] - [[package]] name = "try-lock" version = "0.2.5" diff --git a/examples/rust/Cargo.toml b/examples/rust/Cargo.toml index 8138e6fb2eb8..df441ad669d8 100644 --- a/examples/rust/Cargo.toml +++ b/examples/rust/Cargo.toml @@ -27,14 +27,14 @@ version = "0.2.0" edition = "2021" authors = ["NVIDIA"] license = "Apache-2.0" -homepage = "https://github.com/triton-inference-server/triton_distributed" -repository = "https://github.com/triton-inference-server/triton_distributed" +homepage = "https://github.com/dynemo-ai/dynemo" +repository = "https://github.com/dynemo-ai/dynemo.git" [workspace.dependencies] # local or crates.io -triton-distributed-runtime = { path = "../../lib/runtime" } -triton-distributed-llm = { path = "../../lib/llm" } +dynemo-runtime = { path = "../../lib/runtime" } +dynemo-llm = { path = "../../lib/llm" } # crates.io anyhow = { version = "1" } diff --git a/examples/rust/hello_world/Cargo.toml b/examples/rust/hello_world/Cargo.toml index 050231ee1814..bb51eee5ce00 100644 --- a/examples/rust/hello_world/Cargo.toml +++ b/examples/rust/hello_world/Cargo.toml @@ -22,6 +22,6 @@ license.workspace = true homepage.workspace = true [dependencies] -triton-distributed-runtime = { workspace = true } +dynemo-runtime = { workspace = true } # third-party diff --git a/examples/rust/hello_world/src/bin/client.rs b/examples/rust/hello_world/src/bin/client.rs index bcd829fd5b5b..9422cea9b88a 100644 --- a/examples/rust/hello_world/src/bin/client.rs +++ b/examples/rust/hello_world/src/bin/client.rs @@ -13,11 +13,11 @@ // See the License for the specific language governing permissions and // limitations under the License. -use hello_world::DEFAULT_NAMESPACE; -use triton_distributed_runtime::{ +use dynemo_runtime::{ logging, protocols::annotated::Annotated, stream::StreamExt, DistributedRuntime, Result, Runtime, Worker, }; +use hello_world::DEFAULT_NAMESPACE; fn main() -> Result<()> { logging::init(); diff --git a/examples/rust/hello_world/src/bin/server.rs b/examples/rust/hello_world/src/bin/server.rs index 839905f1d9a8..8d9d1ca983d8 100644 --- a/examples/rust/hello_world/src/bin/server.rs +++ b/examples/rust/hello_world/src/bin/server.rs @@ -13,9 +13,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use hello_world::DEFAULT_NAMESPACE; -use std::sync::Arc; -use triton_distributed_runtime::{ +use dynemo_runtime::{ logging, pipeline::{ async_trait, network::Ingress, AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, @@ -24,6 +22,8 @@ use triton_distributed_runtime::{ protocols::annotated::Annotated, stream, DistributedRuntime, Result, Runtime, Worker, }; +use hello_world::DEFAULT_NAMESPACE; +use std::sync::Arc; fn main() -> Result<()> { logging::init(); diff --git a/examples/rust/hello_world/src/lib.rs b/examples/rust/hello_world/src/lib.rs index 72d0d669be89..5fa0e844a7cc 100644 --- a/examples/rust/hello_world/src/lib.rs +++ b/examples/rust/hello_world/src/lib.rs @@ -13,4 +13,4 @@ // See the License for the specific language governing permissions and // limitations under the License. -pub const DEFAULT_NAMESPACE: &str = "triton-init"; +pub const DEFAULT_NAMESPACE: &str = "dynemo"; diff --git a/examples/rust/http/Cargo.toml b/examples/rust/http/Cargo.toml index 4ed8a0e17f72..5b063c7b5fbb 100644 --- a/examples/rust/http/Cargo.toml +++ b/examples/rust/http/Cargo.toml @@ -24,8 +24,8 @@ homepage.workspace = true repository.workspace = true [dependencies] -triton-distributed-runtime = { workspace = true} -triton-distributed-llm = { workspace = true} +dynemo-runtime = { workspace = true} +dynemo-llm = { workspace = true} clap = { version = "4.5", features = ["derive"] } serde = { workspace = true } diff --git a/examples/rust/http/src/main.rs b/examples/rust/http/src/main.rs index 9319fa4a6d47..683615d8bb6d 100644 --- a/examples/rust/http/src/main.rs +++ b/examples/rust/http/src/main.rs @@ -16,14 +16,14 @@ use clap::Parser; use std::sync::Arc; -use triton_distributed_llm::{ +use dynemo_llm::{ http::service::{ discovery::{model_watcher, ModelWatchState}, service_v2::HttpService, }, model_type::ModelType, }; -use triton_distributed_runtime::{ +use dynemo_runtime::{ logging, transports::etcd::PrefixWatcher, DistributedRuntime, Result, Runtime, Worker, }; diff --git a/examples/rust/llmctl/Cargo.toml b/examples/rust/llmctl/Cargo.toml index e5c502fd6d34..c8ad45bd83aa 100644 --- a/examples/rust/llmctl/Cargo.toml +++ b/examples/rust/llmctl/Cargo.toml @@ -23,8 +23,8 @@ homepage.workspace = true repository.workspace = true [dependencies] -triton-distributed-runtime = { workspace = true} -triton-distributed-llm = { workspace = true} +dynemo-runtime = { workspace = true} +dynemo-llm = { workspace = true} serde = { workspace = true } serde_json = { workspace = true } diff --git a/examples/rust/llmctl/src/main.rs b/examples/rust/llmctl/src/main.rs index 35563dda1073..6d2ef90d0917 100644 --- a/examples/rust/llmctl/src/main.rs +++ b/examples/rust/llmctl/src/main.rs @@ -16,8 +16,8 @@ use clap::{Parser, Subcommand}; use tracing as log; -use triton_distributed_llm::{http::service::discovery::ModelEntry, model_type::ModelType}; -use triton_distributed_runtime::{ +use dynemo_llm::{http::service::discovery::ModelEntry, model_type::ModelType}; +use dynemo_runtime::{ distributed::DistributedConfig, logging, protocols::Endpoint, raise, DistributedRuntime, Result, Runtime, Worker, }; diff --git a/examples/rust/service_metrics/Cargo.toml b/examples/rust/service_metrics/Cargo.toml index 1fe20f5cf529..a9e8c752e80d 100644 --- a/examples/rust/service_metrics/Cargo.toml +++ b/examples/rust/service_metrics/Cargo.toml @@ -23,10 +23,10 @@ homepage.workspace = true repository.workspace = true [dependencies] -triton-distributed-runtime = { workspace = true } +dynemo-runtime = { workspace = true } # third-party futures = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } -tokio = { workspace = true } \ No newline at end of file +tokio = { workspace = true } diff --git a/examples/rust/service_metrics/README.md b/examples/rust/service_metrics/README.md index 4d1139d5f255..d76ca4d1d3fa 100644 --- a/examples/rust/service_metrics/README.md +++ b/examples/rust/service_metrics/README.md @@ -4,14 +4,14 @@ This example extends the hello_world example by calling the `scrape_service` met with the service name for the request response the client just issued a request. ```bash -TRD_LOG=debug cargo run --bin server +DYN_LOG=debug cargo run --bin server ``` The client can now observe some basic statistics about each instance of the service begin hosted. ```bash -TRD_LOG=info cargo run --bin client +DYN_LOG=info cargo run --bin client ``` ## Example Output @@ -27,7 +27,7 @@ Annotated { data: Some("o"), id: None, event: None, comment: None } Annotated { data: Some("r"), id: None, event: None, comment: None } Annotated { data: Some("l"), id: None, event: None, comment: None } Annotated { data: Some("d"), id: None, event: None, comment: None } -ServiceSet { services: [ServiceInfo { name: "triton_init_backend_720278f8", id: "eOHMc4ndRw8s5flv4WOZx7", version: "0.0.1", started: "2025-02-26T18:54:04.917294605Z", endpoints: [EndpointInfo { name: "triton_init_backend_720278f8-generate-694d951a80e06abf", subject: "triton_init_backend_720278f8.generate-694d951a80e06abf", data: Some(Metrics(Object {"average_processing_time": Number(53662), "data": Object {"val": Number(10)}, "last_error": String(""), "num_errors": Number(0), "num_requests": Number(2), "processing_time": Number(107325), "queue_group": String("q")})) }] }] } +ServiceSet { services: [ServiceInfo { name: "dynemo_init_backend_720278f8", id: "eOHMc4ndRw8s5flv4WOZx7", version: "0.0.1", started: "2025-02-26T18:54:04.917294605Z", endpoints: [EndpointInfo { name: "dynemo_init_backend_720278f8-generate-694d951a80e06abf", subject: "dynemo_init_backend_720278f8.generate-694d951a80e06abf", data: Some(Metrics(Object {"average_processing_time": Number(53662), "data": Object {"val": Number(10)}, "last_error": String(""), "num_errors": Number(0), "num_requests": Number(2), "processing_time": Number(107325), "queue_group": String("q")})) }] }] } ``` Note the following stats in the output demonstrate the custom diff --git a/examples/rust/service_metrics/src/bin/client.rs b/examples/rust/service_metrics/src/bin/client.rs index 8b1a4208982c..2f4aedf12564 100644 --- a/examples/rust/service_metrics/src/bin/client.rs +++ b/examples/rust/service_metrics/src/bin/client.rs @@ -16,7 +16,7 @@ use futures::StreamExt; use service_metrics::DEFAULT_NAMESPACE; -use triton_distributed_runtime::{ +use dynemo_runtime::{ logging, protocols::annotated::Annotated, utils::Duration, DistributedRuntime, Result, Runtime, Worker, }; diff --git a/examples/rust/service_metrics/src/bin/server.rs b/examples/rust/service_metrics/src/bin/server.rs index 390af01f51d0..8c50b3067222 100644 --- a/examples/rust/service_metrics/src/bin/server.rs +++ b/examples/rust/service_metrics/src/bin/server.rs @@ -15,8 +15,7 @@ use service_metrics::{MyStats, DEFAULT_NAMESPACE}; -use std::sync::Arc; -use triton_distributed_runtime::{ +use dynemo_runtime::{ logging, pipeline::{ async_trait, network::Ingress, AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, @@ -25,6 +24,7 @@ use triton_distributed_runtime::{ protocols::annotated::Annotated, stream, DistributedRuntime, Result, Runtime, Worker, }; +use std::sync::Arc; fn main() -> Result<()> { logging::init(); diff --git a/examples/rust/service_metrics/src/lib.rs b/examples/rust/service_metrics/src/lib.rs index c69cdca65f01..7e12e7ac1ecc 100644 --- a/examples/rust/service_metrics/src/lib.rs +++ b/examples/rust/service_metrics/src/lib.rs @@ -15,7 +15,7 @@ use serde::{Deserialize, Serialize}; -pub const DEFAULT_NAMESPACE: &str = "triton-init"; +pub const DEFAULT_NAMESPACE: &str = "dynemo"; #[derive(Serialize, Deserialize)] // Dummy Stats object to demonstrate how to attach a custom stats handler diff --git a/launch/dynemo-run/Cargo.lock b/launch/dynemo-run/Cargo.lock index 220b11b7e011..013a3c5b95b4 100644 --- a/launch/dynemo-run/Cargo.lock +++ b/launch/dynemo-run/Cargo.lock @@ -1374,6 +1374,58 @@ dependencies = [ "reborrow", ] +[[package]] +name = "dynemo-llm" +version = "0.2.1" +dependencies = [ + "anyhow", + "async-openai", + "async-stream", + "async-trait", + "async_zmq", + "axum 0.8.1", + "bindgen 0.70.1", + "blake3", + "bs62", + "bytes", + "chrono", + "cmake", + "derive_builder", + "dynemo-runtime", + "either", + "erased-serde", + "futures", + "galil-seiferas", + "indexmap 2.7.1", + "itertools 0.14.0", + "libc", + "llama-cpp-2", + "minijinja", + "minijinja-contrib", + "mistralrs", + "prometheus", + "pyo3", + "regex", + "semver", + "serde", + "serde-pickle", + "serde_json", + "serde_repr", + "strum 0.27.1", + "thiserror 2.0.11", + "tokenizers", + "tokio", + "tokio-stream", + "tokio-util", + "toktrie 0.6.28", + "toktrie_hf_tokenizers 0.6.28", + "tracing", + "unicode-segmentation", + "uuid 1.14.0", + "validator", + "xxhash-rust", +] + [[package]] name = "dynemo-run" version = "0.1.0" @@ -1384,6 +1436,8 @@ dependencies = [ "async-trait", "clap", "dialoguer", + "dynemo-llm", + "dynemo-runtime", "futures", "futures-util", "libc", @@ -1395,8 +1449,50 @@ dependencies = [ "tokio-util", "tracing", "tracing-subscriber", - "triton-distributed-llm", - "triton-distributed-runtime", +] + +[[package]] +name = "dynemo-runtime" +version = "0.2.1" +dependencies = [ + "anyhow", + "async-nats", + "async-once-cell", + "async-stream", + "async-trait", + "async_zmq", + "blake3", + "bytes", + "chrono", + "derive-getters", + "derive_builder", + "educe", + "either", + "etcd-client", + "figment", + "futures", + "humantime", + "local-ip-address", + "log", + "nid", + "nix 0.29.0", + "nuid", + "once_cell", + "prometheus", + "rand", + "regex", + "serde", + "serde_json", + "socket2", + "thiserror 1.0.69", + "tokio", + "tokio-stream", + "tokio-util", + "tracing", + "tracing-subscriber", + "uuid 1.14.0", + "validator", + "xxhash-rust", ] [[package]] @@ -5594,102 +5690,6 @@ dependencies = [ "tracing-serde", ] -[[package]] -name = "triton-distributed-llm" -version = "0.2.1" -dependencies = [ - "anyhow", - "async-openai", - "async-stream", - "async-trait", - "async_zmq", - "axum 0.8.1", - "bindgen 0.70.1", - "blake3", - "bs62", - "bytes", - "chrono", - "cmake", - "derive_builder", - "either", - "erased-serde", - "futures", - "galil-seiferas", - "indexmap 2.7.1", - "itertools 0.14.0", - "libc", - "llama-cpp-2", - "minijinja", - "minijinja-contrib", - "mistralrs", - "prometheus", - "pyo3", - "regex", - "semver", - "serde", - "serde-pickle", - "serde_json", - "serde_repr", - "strum 0.27.1", - "thiserror 2.0.11", - "tokenizers", - "tokio", - "tokio-stream", - "tokio-util", - "toktrie 0.6.28", - "toktrie_hf_tokenizers 0.6.28", - "tracing", - "triton-distributed-runtime", - "unicode-segmentation", - "uuid 1.14.0", - "validator", - "xxhash-rust", -] - -[[package]] -name = "triton-distributed-runtime" -version = "0.2.1" -dependencies = [ - "anyhow", - "async-nats", - "async-once-cell", - "async-stream", - "async-trait", - "async_zmq", - "blake3", - "bytes", - "chrono", - "derive-getters", - "derive_builder", - "educe", - "either", - "etcd-client", - "figment", - "futures", - "humantime", - "local-ip-address", - "log", - "nid", - "nix 0.29.0", - "nuid", - "once_cell", - "prometheus", - "rand", - "regex", - "serde", - "serde_json", - "socket2", - "thiserror 1.0.69", - "tokio", - "tokio-stream", - "tokio-util", - "tracing", - "tracing-subscriber", - "uuid 1.14.0", - "validator", - "xxhash-rust", -] - [[package]] name = "try-lock" version = "0.2.5" diff --git a/launch/dynemo-run/Cargo.toml b/launch/dynemo-run/Cargo.toml index 9d81c344d8c1..dc5c9beb8d4f 100644 --- a/launch/dynemo-run/Cargo.toml +++ b/launch/dynemo-run/Cargo.toml @@ -22,13 +22,13 @@ homepage = "https://github.com/dynemo-ai/dynemo" license = "Apache-2.0" [features] -mistralrs = ["triton-distributed-llm/mistralrs"] -sglang = ["triton-distributed-llm/sglang", "dep:netlink-packet-route", "dep:rtnetlink"] -vllm = ["triton-distributed-llm/vllm", "dep:netlink-packet-route", "dep:rtnetlink"] -llamacpp = ["triton-distributed-llm/llamacpp"] -trtllm = ["triton-distributed-llm/trtllm"] -cuda = ["triton-distributed-llm/cuda"] -metal = ["triton-distributed-llm/metal"] +mistralrs = ["dynemo-llm/mistralrs"] +sglang = ["dynemo-llm/sglang", "dep:netlink-packet-route", "dep:rtnetlink"] +vllm = ["dynemo-llm/vllm", "dep:netlink-packet-route", "dep:rtnetlink"] +llamacpp = ["dynemo-llm/llamacpp"] +trtllm = ["dynemo-llm/trtllm"] +cuda = ["dynemo-llm/cuda"] +metal = ["dynemo-llm/metal"] [dependencies] anyhow = "1" @@ -48,5 +48,5 @@ tokio = { version = "1", features = ["full"] } tokio-util = { version = "0.7", features = ["codec", "net"] } tracing = { version = "0.1" } tracing-subscriber = { version = "0.3", features = ["env-filter", "local-time", "json"] } -triton-distributed-runtime = { path = "../../lib/runtime" } -triton-distributed-llm = { path = "../../lib/llm" } +dynemo-runtime = { path = "../../lib/runtime" } +dynemo-llm = { path = "../../lib/llm" } diff --git a/launch/dynemo-run/src/input/endpoint.rs b/launch/dynemo-run/src/input/endpoint.rs index ae3bc26d67cf..45d85748da50 100644 --- a/launch/dynemo-run/src/input/endpoint.rs +++ b/launch/dynemo-run/src/input/endpoint.rs @@ -13,7 +13,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use triton_distributed_llm::{ +use dynemo_llm::{ backend::Backend, http::service::discovery::ModelEntry, model_type::ModelType, @@ -25,10 +25,10 @@ use triton_distributed_llm::{ Annotated, }, }; -use triton_distributed_runtime::pipeline::{ +use dynemo_runtime::pipeline::{ network::Ingress, ManyOut, Operator, SegmentSource, ServiceBackend, SingleIn, Source, }; -use triton_distributed_runtime::{protocols::Endpoint, DistributedRuntime, Runtime}; +use dynemo_runtime::{protocols::Endpoint, DistributedRuntime, Runtime}; use crate::EngineConfig; diff --git a/launch/dynemo-run/src/input/http.rs b/launch/dynemo-run/src/input/http.rs index 98d9e1bf051f..9a01f4b35865 100644 --- a/launch/dynemo-run/src/input/http.rs +++ b/launch/dynemo-run/src/input/http.rs @@ -15,7 +15,7 @@ use std::sync::Arc; -use triton_distributed_llm::{ +use dynemo_llm::{ backend::Backend, http::service::{discovery, service_v2}, model_type::ModelType, @@ -27,7 +27,7 @@ use triton_distributed_llm::{ Annotated, }, }; -use triton_distributed_runtime::{ +use dynemo_runtime::{ pipeline::{ManyOut, Operator, ServiceBackend, ServiceFrontend, SingleIn, Source}, DistributedRuntime, Runtime, }; diff --git a/launch/dynemo-run/src/input/text.rs b/launch/dynemo-run/src/input/text.rs index a9c0f9267d33..a838f9128c8c 100644 --- a/launch/dynemo-run/src/input/text.rs +++ b/launch/dynemo-run/src/input/text.rs @@ -13,12 +13,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use futures::StreamExt; -use std::{ - io::{ErrorKind, Read, Write}, - sync::Arc, -}; -use triton_distributed_llm::{ +use dynemo_llm::{ backend::Backend, preprocessor::OpenAIPreprocessor, types::{ @@ -29,10 +24,15 @@ use triton_distributed_llm::{ Annotated, }, }; -use triton_distributed_runtime::{ +use dynemo_runtime::{ pipeline::{Context, ManyOut, Operator, ServiceBackend, ServiceFrontend, SingleIn, Source}, runtime::CancellationToken, }; +use futures::StreamExt; +use std::{ + io::{ErrorKind, Read, Write}, + sync::Arc, +}; use crate::EngineConfig; diff --git a/launch/dynemo-run/src/lib.rs b/launch/dynemo-run/src/lib.rs index e7566041c087..74a83a52f960 100644 --- a/launch/dynemo-run/src/lib.rs +++ b/launch/dynemo-run/src/lib.rs @@ -16,7 +16,7 @@ #[cfg(any(feature = "vllm", feature = "sglang"))] use std::{future::Future, pin::Pin}; -use triton_distributed_llm::{ +use dynemo_llm::{ backend::ExecutionContext, model_card::model::ModelDeploymentCard, types::{ @@ -27,7 +27,7 @@ use triton_distributed_llm::{ Annotated, }, }; -use triton_distributed_runtime::{component::Client, protocols::Endpoint, DistributedRuntime}; +use dynemo_runtime::{component::Client, protocols::Endpoint, DistributedRuntime}; mod flags; pub use flags::Flags; @@ -67,7 +67,7 @@ pub enum EngineConfig { #[allow(unused_mut)] pub async fn run( - runtime: triton_distributed_runtime::Runtime, + runtime: dynemo_runtime::Runtime, mut in_opt: Input, // mut because vllm and sglang multi-node can change it out_opt: Output, flags: Flags, @@ -173,13 +173,12 @@ pub async fn run( }; EngineConfig::StaticFull { service_name: model_name, - engine: triton_distributed_llm::engines::mistralrs::make_engine(&model_path) - .await?, + engine: dynemo_llm::engines::mistralrs::make_engine(&model_path).await?, } } #[cfg(feature = "sglang")] Output::SgLang => { - use triton_distributed_llm::engines::sglang; + use dynemo_llm::engines::sglang; let Some(model_path) = model_path else { anyhow::bail!("out=sglang requires flag --model-path="); }; @@ -191,7 +190,7 @@ pub async fn run( let Some(sock_prefix) = zmq_socket_prefix else { anyhow::bail!("sglang requires zmq_socket_prefix"); }; - let node_conf = triton_distributed_llm::engines::MultiNodeConfig { + let node_conf = dynemo_llm::engines::MultiNodeConfig { num_nodes: flags.num_nodes, node_rank: flags.node_rank, leader_addr: flags.leader_addr.unwrap_or_default(), @@ -229,7 +228,7 @@ pub async fn run( } #[cfg(feature = "vllm")] Output::Vllm => { - use triton_distributed_llm::engines::vllm; + use dynemo_llm::engines::vllm; if flags.base_gpu_id != 0 { anyhow::bail!("vllm does not support base_gpu_id. Set environment variable CUDA_VISIBLE_DEVICES instead."); } @@ -253,7 +252,7 @@ pub async fn run( let Some(sock_prefix) = zmq_socket_prefix else { anyhow::bail!("vllm requires zmq_socket_prefix"); }; - let node_conf = triton_distributed_llm::engines::MultiNodeConfig { + let node_conf = dynemo_llm::engines::MultiNodeConfig { num_nodes: flags.num_nodes, node_rank: flags.node_rank, leader_addr: flags.leader_addr.unwrap_or_default(), @@ -296,7 +295,7 @@ pub async fn run( } #[cfg(feature = "llamacpp")] Output::LlamaCpp => { - use triton_distributed_llm::engines::llamacpp; + use dynemo_llm::engines::llamacpp; let Some(model_path) = model_path else { anyhow::bail!("out=llamacpp requires flag --model-path="); }; @@ -317,7 +316,7 @@ pub async fn run( } #[cfg(feature = "trtllm")] Output::TrtLLM => { - use triton_distributed_llm::engines::trtllm; + use dynemo_llm::engines::trtllm; let Some(model_path) = model_path else { anyhow::bail!("out=trtllm requires flag --model-path="); }; diff --git a/launch/dynemo-run/src/main.rs b/launch/dynemo-run/src/main.rs index 2b4096c4898c..3c8a48cbf405 100644 --- a/launch/dynemo-run/src/main.rs +++ b/launch/dynemo-run/src/main.rs @@ -18,7 +18,7 @@ use std::env; use clap::Parser; use dynemo_run::{Input, Output}; -use triton_distributed_runtime::logging; +use dynemo_runtime::logging; const HELP: &str = r#" dynemo-run is a single binary that wires together the various inputs (http, text, network) and workers (network, engine), that runs the services. It is the simplest way to use dynemo locally. @@ -60,13 +60,13 @@ fn main() -> anyhow::Result<()> { if cfg!(feature = "sglang") { #[cfg(feature = "sglang")] { - use triton_distributed_llm::engines::sglang; + use dynemo_llm::engines::sglang; let gpu_config = sglang::MultiGPUConfig { tp_size: flags.tensor_parallel_size, tp_rank: sglang_flags.tp_rank, gpu_id: sglang_flags.gpu_id, }; - let node_config = triton_distributed_llm::engines::MultiNodeConfig { + let node_config = dynemo_llm::engines::MultiNodeConfig { num_nodes: flags.num_nodes, node_rank: flags.node_rank, leader_addr: flags.leader_addr.unwrap_or_default(), @@ -98,8 +98,8 @@ fn main() -> anyhow::Result<()> { if cfg!(feature = "vllm") { #[cfg(feature = "vllm")] { - use triton_distributed_llm::engines::vllm; - let node_config = triton_distributed_llm::engines::MultiNodeConfig { + use dynemo_llm::engines::vllm; + let node_config = dynemo_llm::engines::MultiNodeConfig { num_nodes: flags.num_nodes, node_rank: flags.node_rank, leader_addr: flags.leader_addr.unwrap_or_default(), @@ -119,15 +119,15 @@ fn main() -> anyhow::Result<()> { } // max_worker_threads and max_blocking_threads from env vars or config file. - let rt_config = triton_distributed_runtime::RuntimeConfig::from_settings()?; + let rt_config = dynemo_runtime::RuntimeConfig::from_settings()?; // One per process. Wraps a Runtime with holds two tokio runtimes. - let worker = triton_distributed_runtime::Worker::from_config(rt_config)?; + let worker = dynemo_runtime::Worker::from_config(rt_config)?; worker.execute(wrapper) } -async fn wrapper(runtime: triton_distributed_runtime::Runtime) -> anyhow::Result<()> { +async fn wrapper(runtime: dynemo_runtime::Runtime) -> anyhow::Result<()> { let mut in_opt = None; let mut out_opt = None; let args: Vec = env::args().skip(1).collect(); diff --git a/launch/dynemo-run/src/output/echo_core.rs b/launch/dynemo-run/src/output/echo_core.rs index 2218a509c876..1a5b7ac11d75 100644 --- a/launch/dynemo-run/src/output/echo_core.rs +++ b/launch/dynemo-run/src/output/echo_core.rs @@ -18,12 +18,12 @@ use std::{sync::Arc, time::Duration}; use async_stream::stream; use async_trait::async_trait; -use triton_distributed_llm::backend::ExecutionContext; -use triton_distributed_llm::preprocessor::BackendInput; -use triton_distributed_llm::protocols::common::llm_backend::LLMEngineOutput; -use triton_distributed_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream}; -use triton_distributed_runtime::pipeline::{Error, ManyOut, SingleIn}; -use triton_distributed_runtime::protocols::annotated::Annotated; +use dynemo_llm::backend::ExecutionContext; +use dynemo_llm::preprocessor::BackendInput; +use dynemo_llm::protocols::common::llm_backend::LLMEngineOutput; +use dynemo_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream}; +use dynemo_runtime::pipeline::{Error, ManyOut, SingleIn}; +use dynemo_runtime::protocols::annotated::Annotated; /// How long to sleep between echoed tokens. /// 50ms gives us 20 tok/s. diff --git a/launch/dynemo-run/src/output/echo_full.rs b/launch/dynemo-run/src/output/echo_full.rs index 64e092fd6078..9099e53cc60d 100644 --- a/launch/dynemo-run/src/output/echo_full.rs +++ b/launch/dynemo-run/src/output/echo_full.rs @@ -18,13 +18,13 @@ use std::{sync::Arc, time::Duration}; use async_stream::stream; use async_trait::async_trait; -use triton_distributed_llm::protocols::openai::chat_completions::{ +use dynemo_llm::protocols::openai::chat_completions::{ NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse, }; -use triton_distributed_llm::types::openai::chat_completions::OpenAIChatCompletionsStreamingEngine; -use triton_distributed_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream}; -use triton_distributed_runtime::pipeline::{Error, ManyOut, SingleIn}; -use triton_distributed_runtime::protocols::annotated::Annotated; +use dynemo_llm::types::openai::chat_completions::OpenAIChatCompletionsStreamingEngine; +use dynemo_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream}; +use dynemo_runtime::pipeline::{Error, ManyOut, SingleIn}; +use dynemo_runtime::protocols::annotated::Annotated; /// How long to sleep between echoed tokens. /// 50ms gives us 20 tok/s. diff --git a/lib/bindings/c/Cargo.lock b/lib/bindings/c/Cargo.lock index fe57528584a8..25bf8bc8b81e 100644 --- a/lib/bindings/c/Cargo.lock +++ b/lib/bindings/c/Cargo.lock @@ -954,6 +954,99 @@ dependencies = [ "syn 2.0.96", ] +[[package]] +name = "dynemo-llm" +version = "0.2.1" +dependencies = [ + "anyhow", + "async-openai", + "async-stream", + "async-trait", + "axum 0.8.1", + "bindgen", + "blake3", + "bs62", + "bytes", + "chrono", + "cmake", + "derive_builder", + "dynemo-runtime", + "either", + "erased-serde", + "futures", + "galil-seiferas", + "indexmap 2.7.0", + "itertools 0.14.0", + "libc", + "minijinja", + "minijinja-contrib", + "prometheus", + "pyo3", + "regex", + "semver", + "serde", + "serde-pickle", + "serde_json", + "serde_repr", + "strum", + "thiserror 2.0.11", + "tokenizers", + "tokio", + "tokio-stream", + "tokio-util", + "toktrie", + "toktrie_hf_tokenizers", + "tracing", + "unicode-segmentation", + "uuid", + "validator", + "xxhash-rust", +] + +[[package]] +name = "dynemo-runtime" +version = "0.2.1" +dependencies = [ + "anyhow", + "async-nats", + "async-once-cell", + "async-stream", + "async-trait", + "async_zmq", + "blake3", + "bytes", + "chrono", + "derive-getters", + "derive_builder", + "educe", + "either", + "etcd-client", + "figment", + "futures", + "humantime", + "local-ip-address", + "log", + "nid", + "nix", + "nuid", + "once_cell", + "prometheus", + "rand", + "regex", + "serde", + "serde_json", + "socket2", + "thiserror 1.0.69", + "tokio", + "tokio-stream", + "tokio-util", + "tracing", + "tracing-subscriber", + "uuid", + "validator", + "xxhash-rust", +] + [[package]] name = "ed25519" version = "2.2.3" @@ -1853,6 +1946,27 @@ version = "0.2.169" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b5aba8db14291edd000dfcc4d620c7ebfb122c613afb886ca8803fa4e128a20a" +[[package]] +name = "libdynemo-llm" +version = "0.2.1" +dependencies = [ + "anyhow", + "async-once-cell", + "cbindgen", + "dynemo-llm", + "dynemo-runtime", + "futures", + "libc", + "once_cell", + "serde", + "serde_json", + "tokio", + "tokio-stream", + "tracing", + "tracing-subscriber", + "uuid", +] + [[package]] name = "libloading" version = "0.8.6" @@ -1873,27 +1987,6 @@ dependencies = [ "libc", ] -[[package]] -name = "libtriton-distributed-llm" -version = "0.2.1" -dependencies = [ - "anyhow", - "async-once-cell", - "cbindgen", - "futures", - "libc", - "once_cell", - "serde", - "serde_json", - "tokio", - "tokio-stream", - "tracing", - "tracing-subscriber", - "triton-distributed-llm", - "triton-distributed-runtime", - "uuid", -] - [[package]] name = "linux-raw-sys" version = "0.4.15" @@ -3955,99 +4048,6 @@ dependencies = [ "tracing-serde", ] -[[package]] -name = "triton-distributed-llm" -version = "0.2.1" -dependencies = [ - "anyhow", - "async-openai", - "async-stream", - "async-trait", - "axum 0.8.1", - "bindgen", - "blake3", - "bs62", - "bytes", - "chrono", - "cmake", - "derive_builder", - "either", - "erased-serde", - "futures", - "galil-seiferas", - "indexmap 2.7.0", - "itertools 0.14.0", - "libc", - "minijinja", - "minijinja-contrib", - "prometheus", - "pyo3", - "regex", - "semver", - "serde", - "serde-pickle", - "serde_json", - "serde_repr", - "strum", - "thiserror 2.0.11", - "tokenizers", - "tokio", - "tokio-stream", - "tokio-util", - "toktrie", - "toktrie_hf_tokenizers", - "tracing", - "triton-distributed-runtime", - "unicode-segmentation", - "uuid", - "validator", - "xxhash-rust", -] - -[[package]] -name = "triton-distributed-runtime" -version = "0.2.1" -dependencies = [ - "anyhow", - "async-nats", - "async-once-cell", - "async-stream", - "async-trait", - "async_zmq", - "blake3", - "bytes", - "chrono", - "derive-getters", - "derive_builder", - "educe", - "either", - "etcd-client", - "figment", - "futures", - "humantime", - "local-ip-address", - "log", - "nid", - "nix", - "nuid", - "once_cell", - "prometheus", - "rand", - "regex", - "serde", - "serde_json", - "socket2", - "thiserror 1.0.69", - "tokio", - "tokio-stream", - "tokio-util", - "tracing", - "tracing-subscriber", - "uuid", - "validator", - "xxhash-rust", -] - [[package]] name = "try-lock" version = "0.2.5" diff --git a/lib/bindings/c/Cargo.toml b/lib/bindings/c/Cargo.toml index 7e1939948f7f..6fe05d5122f5 100644 --- a/lib/bindings/c/Cargo.toml +++ b/lib/bindings/c/Cargo.toml @@ -14,24 +14,24 @@ # limitations under the License. [package] -name = "libtriton-distributed-llm" +name = "libdynemo-llm" version = "0.2.1" edition = "2021" authors = ["NVIDIA"] license = "Apache-2.0" -homepage = "https://github.com/triton-inference-server/triton_distributed" -repository = "https://github.com/triton-inference-server/triton_distributed" +homepage = "https://github.com/dynemo-ai/dynemo" +repository = "https://github.com/dynemo-ai/dynemo.git" [lib] -name = "triton_distributed_llm_capi" +name = "dynemo_llm_capi" crate-type = ["cdylib"] [build-dependencies] cbindgen = "0.27" [dependencies] -triton-distributed-llm = { path = "../../llm" } -triton-distributed-runtime = { path = "../../runtime" } +dynemo-llm = { path = "../../llm" } +dynemo-runtime = { path = "../../runtime" } anyhow = { version = "1" } futures = "0.3" diff --git a/lib/bindings/c/build.rs b/lib/bindings/c/build.rs index 4a692ce8b340..0f11b97084e1 100644 --- a/lib/bindings/c/build.rs +++ b/lib/bindings/c/build.rs @@ -22,7 +22,7 @@ fn main() { let header_path = Path::new(&crate_dir) .join("include") .join("nvidia") - .join("triton_llm") + .join("dynemo_llm") .join("llm_engine.h"); cbindgen::generate(crate_dir) diff --git a/lib/bindings/c/cbindgen.toml b/lib/bindings/c/cbindgen.toml index 4cb4b824df42..582f1cb0f51c 100644 --- a/lib/bindings/c/cbindgen.toml +++ b/lib/bindings/c/cbindgen.toml @@ -15,7 +15,7 @@ language = "C++" cpp_compat = true -include_guard = "__NVIDIA_TRITON_LLM_API__" +include_guard = "__NVIDIA_DYNEMO_LLM_API__" [enum] @@ -25,7 +25,7 @@ enum_class = false [export] -include = ["TritonLlmResult", "triton_llm_init", "triton_llm_shutdown"] +include = ["DynemoLlmResult", "dynemo_llm_init", "dynemo_llm_shutdown"] [export.rename] -"TritonLlmResult" = "triton_llm_result_t" +"DynemoLlmResult" = "dynemo_llm_result_t" diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index 35c59d9b7a41..3d0be32239d1 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -19,10 +19,10 @@ use once_cell::sync::OnceCell; use std::ffi::CStr; use std::sync::atomic::{AtomicU32, Ordering}; -use triton_distributed_llm::kv_router::{ +use dynemo_llm::kv_router::{ indexer::compute_block_hash_for_seq, protocols::*, publisher::KvEventPublisher, }; -use triton_distributed_runtime::{DistributedRuntime, Worker}; +use dynemo_runtime::{DistributedRuntime, Worker}; static WK: OnceCell = OnceCell::new(); static DRT: AsyncOnceCell = AsyncOnceCell::new(); // [FIXME] shouldn't the publisher be instance passing between API calls? @@ -41,7 +41,7 @@ fn initialize_tracing() { } #[repr(u32)] -pub enum TritonLlmResult { +pub enum DynemoLlmResult { OK = 0, ERR = 1, } @@ -49,17 +49,17 @@ pub enum TritonLlmResult { /// # Safety /// the namespace_c_str and component_c_str are passed as pointers to C strings #[no_mangle] -pub unsafe extern "C" fn triton_llm_init( +pub unsafe extern "C" fn dynemo_llm_init( namespace_c_str: *const c_char, component_c_str: *const c_char, worker_id: i64, -) -> TritonLlmResult { +) -> DynemoLlmResult { initialize_tracing(); let wk = match WK.get_or_try_init(Worker::from_settings) { Ok(wk) => wk.clone(), Err(e) => { eprintln!("Failed to initialize runtime: {:?}", e); - return TritonLlmResult::ERR; + return DynemoLlmResult::ERR; } }; let rt = wk.runtime(); @@ -73,7 +73,7 @@ pub unsafe extern "C" fn triton_llm_init( Ok(_) => Ok(()), Err(e) => { eprintln!("Failed to initialize distributed runtime: {:?}", e); - Err(TritonLlmResult::ERR) + Err(DynemoLlmResult::ERR) } } }); @@ -81,7 +81,7 @@ pub unsafe extern "C" fn triton_llm_init( Ok(s) => s.to_string(), Err(e) => { eprintln!("Failed to convert C string to Rust string: {:?}", e); - return TritonLlmResult::ERR; + return DynemoLlmResult::ERR; } }; @@ -89,18 +89,18 @@ pub unsafe extern "C" fn triton_llm_init( Ok(s) => s.to_string(), Err(e) => { eprintln!("Failed to convert C string to Rust string: {:?}", e); - return TritonLlmResult::ERR; + return DynemoLlmResult::ERR; } }; match result { Ok(_) => match KV_PUB - .get_or_try_init(move || triton_create_kv_publisher(namespace, component, worker_id)) + .get_or_try_init(move || dynemo_create_kv_publisher(namespace, component, worker_id)) { - Ok(_) => TritonLlmResult::OK, + Ok(_) => DynemoLlmResult::OK, Err(e) => { eprintln!("Failed to initialize distributed runtime: {:?}", e); - TritonLlmResult::ERR + DynemoLlmResult::ERR } }, Err(e) => e, @@ -108,33 +108,33 @@ pub unsafe extern "C" fn triton_llm_init( } #[no_mangle] -pub extern "C" fn triton_llm_shutdown() -> TritonLlmResult { +pub extern "C" fn dynemo_llm_shutdown() -> DynemoLlmResult { let wk = match WK.get() { Some(wk) => wk, None => { eprintln!("Runtime not initialized"); - return TritonLlmResult::ERR; + return DynemoLlmResult::ERR; } }; wk.runtime().shutdown(); - TritonLlmResult::OK + DynemoLlmResult::OK } #[no_mangle] -pub extern "C" fn triton_llm_load_publisher_create() -> TritonLlmResult { - TritonLlmResult::OK +pub extern "C" fn dynemo_llm_load_publisher_create() -> DynemoLlmResult { + DynemoLlmResult::OK } // instantiate a kv publisher // this will bring up the task to publish and the channels to await publishing events -// the [`triton_kv_publish_store_event`] call will use a handle to the publisher to send events -// store and the [`triton_kv_event_create_removed`] will create remove events +// the [`dynemo_kv_publish_store_event`] call will use a handle to the publisher to send events +// store and the [`dynemo_kv_event_create_removed`] will create remove events // these call mus be driving by external c++ threads that are consuming the kv events from the // c++ executor api -fn triton_create_kv_publisher( +fn dynemo_create_kv_publisher( namespace: String, component: String, worker_id: i64, @@ -238,7 +238,7 @@ fn kv_event_create_removed_from_parts( /// parent_hash is passed as pointer to indicate whether the blocks /// has a parent hash or not. nullptr is used to represent no parent hash #[no_mangle] -pub unsafe extern "C" fn triton_kv_event_publish_stored( +pub unsafe extern "C" fn dynemo_kv_event_publish_stored( event_id: u64, token_ids: *const u32, num_block_tokens: *const usize, @@ -246,7 +246,7 @@ pub unsafe extern "C" fn triton_kv_event_publish_stored( num_blocks: usize, parent_hash: *const u64, lora_id: u64, -) -> TritonLlmResult { +) -> DynemoLlmResult { let publisher = KV_PUB.get().unwrap(); let parent_hash = { if parent_hash.is_null() { @@ -265,40 +265,40 @@ pub unsafe extern "C" fn triton_kv_event_publish_stored( lora_id, ); match publisher.publish(event) { - Ok(_) => TritonLlmResult::OK, + Ok(_) => DynemoLlmResult::OK, Err(e) => { eprintln!("Error publishing stored kv event {:?}", e); - TritonLlmResult::ERR + DynemoLlmResult::ERR } } } #[no_mangle] -pub extern "C" fn triton_kv_event_publish_removed( +pub extern "C" fn dynemo_kv_event_publish_removed( event_id: u64, block_ids: *const u64, num_blocks: usize, -) -> TritonLlmResult { +) -> DynemoLlmResult { let publisher = KV_PUB.get().unwrap(); let event = kv_event_create_removed_from_parts(event_id, block_ids, num_blocks); match publisher.publish(event) { - Ok(_) => TritonLlmResult::OK, + Ok(_) => DynemoLlmResult::OK, Err(e) => { eprintln!("Error publishing removed kv event {:?}", e); - TritonLlmResult::ERR + DynemoLlmResult::ERR } } } // #[no_mangle] -// pub extern "C" fn triton_kv_publish_store_event( +// pub extern "C" fn dynemo_kv_publish_store_event( // event_id: u64, // token_ids: *const u32, // num_tokens: usize, // lora_id: u64, -// ) -> TritonLlmResult { +// ) -> DynemoLlmResult { // // if event.is_null() || token_ids.is_null() { -// // return tritonKvErrorType::INVALID_TOKEN_IDS; +// // return dynemoKvErrorType::INVALID_TOKEN_IDS; // // } // // let tokens = unsafe { std::slice::from_raw_parts(token_ids, num_tokens) }.to_vec(); @@ -311,15 +311,15 @@ pub extern "C" fn triton_kv_event_publish_removed( // // unsafe { *event = Box::into_raw(new_event) }; -// TritonLlmResult::OK +// DynemoLlmResult::OK // } // #[no_mangle] -// pub extern "C" fn triton_kv_event_create_removed( +// pub extern "C" fn dynemo_kv_event_create_removed( // event_id: u64, // block_hashes: *const u64, // num_hashes: usize, -// ) -> TritonLlmResult { +// ) -> DynemoLlmResult { // // if event.is_null() || block_hashes.is_null() { // // return -1; // // } @@ -334,19 +334,19 @@ pub extern "C" fn triton_kv_event_publish_removed( // // unsafe { *event = Box::into_raw(new_event) }; // // 0 -// TritonLlmResult::OK +// DynemoLlmResult::OK // } // /// create load publisher object and return a handle // /// load publisher will instantiate the nats service and tie its stats handler to // /// a watch channel receiver. the watch channel sender will be attach to the -// /// handle and calls to [`triton_load_stats_publish`] issue the stats to the watch t -// pub extern "C" fn triton_load_publisher_create() -> *mut LoadPublisher { +// /// handle and calls to [`dynemo_load_stats_publish`] issue the stats to the watch t +// pub extern "C" fn dynemo_load_publisher_create() -> *mut LoadPublisher { // // let publisher = Box::new(LoadPublisher::new()); // // Box::into_raw(publisher) // } -// pub extern "C" fn triton_load_stats_publish( +// pub extern "C" fn dynemo_load_stats_publish( // publisher: *mut LoadPublisher, // active_slots: u64, // total_slots: u64, diff --git a/lib/bindings/python/.gitignore b/lib/bindings/python/.gitignore index 344ab95885db..eda1176507c0 100644 --- a/lib/bindings/python/.gitignore +++ b/lib/bindings/python/.gitignore @@ -1,2 +1,2 @@ /target -python/triton_distributed/*.so +python/dynemo/.*.so diff --git a/lib/bindings/python/Cargo.lock b/lib/bindings/python/Cargo.lock index f02d32869507..a5d693a27f43 100644 --- a/lib/bindings/python/Cargo.lock +++ b/lib/bindings/python/Cargo.lock @@ -956,6 +956,119 @@ dependencies = [ "syn 2.0.98", ] +[[package]] +name = "dynemo-llm" +version = "0.2.1" +dependencies = [ + "anyhow", + "async-openai", + "async-stream", + "async-trait", + "axum 0.8.1", + "bindgen", + "blake3", + "bs62", + "bytes", + "chrono", + "cmake", + "derive_builder", + "dynemo-runtime", + "either", + "erased-serde", + "futures", + "galil-seiferas", + "indexmap 2.7.1", + "itertools 0.14.0", + "libc", + "minijinja", + "minijinja-contrib", + "prometheus", + "pyo3", + "regex", + "semver", + "serde", + "serde-pickle", + "serde_json", + "serde_repr", + "strum", + "thiserror 2.0.11", + "tokenizers", + "tokio", + "tokio-stream", + "tokio-util", + "toktrie", + "toktrie_hf_tokenizers", + "tracing", + "unicode-segmentation", + "uuid", + "validator", + "xxhash-rust", +] + +[[package]] +name = "dynemo-py3" +version = "0.2.1" +dependencies = [ + "dynemo-llm", + "dynemo-runtime", + "futures", + "once_cell", + "pyo3", + "pyo3-async-runtimes", + "pythonize", + "serde", + "serde_json", + "thiserror 2.0.11", + "tokio", + "tokio-stream", + "tracing", + "tracing-subscriber", +] + +[[package]] +name = "dynemo-runtime" +version = "0.2.1" +dependencies = [ + "anyhow", + "async-nats", + "async-once-cell", + "async-stream", + "async-trait", + "async_zmq", + "blake3", + "bytes", + "chrono", + "derive-getters", + "derive_builder", + "educe", + "either", + "etcd-client", + "figment", + "futures", + "humantime", + "local-ip-address", + "log", + "nid", + "nix", + "nuid", + "once_cell", + "prometheus", + "rand", + "regex", + "serde", + "serde_json", + "socket2", + "thiserror 1.0.69", + "tokio", + "tokio-stream", + "tokio-util", + "tracing", + "tracing-subscriber", + "uuid", + "validator", + "xxhash-rust", +] + [[package]] name = "ed25519" version = "2.2.3" @@ -4004,119 +4117,6 @@ dependencies = [ "tracing-serde", ] -[[package]] -name = "triton-distributed-llm" -version = "0.2.1" -dependencies = [ - "anyhow", - "async-openai", - "async-stream", - "async-trait", - "axum 0.8.1", - "bindgen", - "blake3", - "bs62", - "bytes", - "chrono", - "cmake", - "derive_builder", - "either", - "erased-serde", - "futures", - "galil-seiferas", - "indexmap 2.7.1", - "itertools 0.14.0", - "libc", - "minijinja", - "minijinja-contrib", - "prometheus", - "pyo3", - "regex", - "semver", - "serde", - "serde-pickle", - "serde_json", - "serde_repr", - "strum", - "thiserror 2.0.11", - "tokenizers", - "tokio", - "tokio-stream", - "tokio-util", - "toktrie", - "toktrie_hf_tokenizers", - "tracing", - "triton-distributed-runtime", - "unicode-segmentation", - "uuid", - "validator", - "xxhash-rust", -] - -[[package]] -name = "triton-distributed-py3" -version = "0.2.1" -dependencies = [ - "futures", - "once_cell", - "pyo3", - "pyo3-async-runtimes", - "pythonize", - "serde", - "serde_json", - "thiserror 2.0.11", - "tokio", - "tokio-stream", - "tracing", - "tracing-subscriber", - "triton-distributed-llm", - "triton-distributed-runtime", -] - -[[package]] -name = "triton-distributed-runtime" -version = "0.2.1" -dependencies = [ - "anyhow", - "async-nats", - "async-once-cell", - "async-stream", - "async-trait", - "async_zmq", - "blake3", - "bytes", - "chrono", - "derive-getters", - "derive_builder", - "educe", - "either", - "etcd-client", - "figment", - "futures", - "humantime", - "local-ip-address", - "log", - "nid", - "nix", - "nuid", - "once_cell", - "prometheus", - "rand", - "regex", - "serde", - "serde_json", - "socket2", - "thiserror 1.0.69", - "tokio", - "tokio-stream", - "tokio-util", - "tracing", - "tracing-subscriber", - "uuid", - "validator", - "xxhash-rust", -] - [[package]] name = "try-lock" version = "0.2.5" diff --git a/lib/bindings/python/Cargo.toml b/lib/bindings/python/Cargo.toml index 496eb89953c7..2f3aed613367 100644 --- a/lib/bindings/python/Cargo.toml +++ b/lib/bindings/python/Cargo.toml @@ -14,13 +14,13 @@ # limitations under the License. [package] -name = "triton-distributed-py3" +name = "dynemo-py3" version = "0.2.1" edition = "2021" authors = ["NVIDIA"] license = "Apache-2.0" -homepage = "https://github.com/triton-inference-server/triton_distributed" -repository = "https://github.com/triton-inference-server/triton_distributed" +homepage = "https://github.com/dynemo-ai/dynemo" +repository = "https://github.com/dynemo-ai/dynemo.git" [lib] path = "rust/lib.rs" @@ -30,8 +30,8 @@ crate-type = ["cdylib"] [dependencies] -triton-distributed-llm = { path = "../../llm" } -triton-distributed-runtime = { path = "../../runtime" } +dynemo-llm = { path = "../../llm" } +dynemo-runtime = { path = "../../runtime" } futures = "0.3" once_cell = "1.20.3" diff --git a/lib/bindings/python/README.md b/lib/bindings/python/README.md index f96c5a3e57f4..3a0bcd68e8f4 100644 --- a/lib/bindings/python/README.md +++ b/lib/bindings/python/README.md @@ -41,7 +41,7 @@ source .venv/bin/activate uv pip install maturin ``` -4. Build and install triton_distributed wheel +4. Build and install dynemo wheel ``` maturin develop --uv ``` diff --git a/lib/bindings/python/examples/bls/bar.py b/lib/bindings/python/examples/bls/bar.py index 6d4e308544df..539661f73896 100644 --- a/lib/bindings/python/examples/bls/bar.py +++ b/lib/bindings/python/examples/bls/bar.py @@ -17,7 +17,7 @@ import uvloop -from triton_distributed.runtime import DistributedRuntime, triton_worker +from dynemo.runtime import DistributedRuntime, dynemo_worker uvloop.install() @@ -29,7 +29,7 @@ async def generate(self, request): yield char -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime): component = runtime.namespace("examples/bls").component("bar") await component.create_service() diff --git a/lib/bindings/python/examples/bls/bls.py b/lib/bindings/python/examples/bls/bls.py index ab17e68cb049..a929186bc9a7 100644 --- a/lib/bindings/python/examples/bls/bls.py +++ b/lib/bindings/python/examples/bls/bls.py @@ -17,12 +17,12 @@ import uvloop -from triton_distributed.runtime import DistributedRuntime, triton_worker +from dynemo.runtime import DistributedRuntime, dynemo_worker uvloop.install() -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime): foo = ( await runtime.namespace("examples/bls") diff --git a/lib/bindings/python/examples/bls/foo.py b/lib/bindings/python/examples/bls/foo.py index e944a136bf5b..4a4d32664a42 100644 --- a/lib/bindings/python/examples/bls/foo.py +++ b/lib/bindings/python/examples/bls/foo.py @@ -17,7 +17,7 @@ import uvloop -from triton_distributed.runtime import DistributedRuntime, triton_worker +from dynemo.runtime import DistributedRuntime, dynemo_worker uvloop.install() @@ -28,7 +28,7 @@ async def generate(self, request): yield char -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime): component = runtime.namespace("examples/bls").component("foo") await component.create_service() diff --git a/lib/bindings/python/examples/error_handling/client.py b/lib/bindings/python/examples/error_handling/client.py index 19f8633d1695..44b9526ba3b0 100644 --- a/lib/bindings/python/examples/error_handling/client.py +++ b/lib/bindings/python/examples/error_handling/client.py @@ -17,12 +17,12 @@ import uvloop -from triton_distributed.runtime import DistributedRuntime, triton_worker +from dynemo.runtime import DistributedRuntime, dynemo_worker -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime): - await init(runtime, "triton-init") + await init(runtime, "dynemo") async def init(runtime: DistributedRuntime, ns: str): diff --git a/lib/bindings/python/examples/error_handling/run.py b/lib/bindings/python/examples/error_handling/run.py index 0f53a73b959f..b73d6011c1d9 100644 --- a/lib/bindings/python/examples/error_handling/run.py +++ b/lib/bindings/python/examples/error_handling/run.py @@ -21,7 +21,7 @@ from client import init as client_init from server import init as server_init -from triton_distributed.runtime import DistributedRuntime, triton_worker +from dynemo.runtime import DistributedRuntime, dynemo_worker def random_string(length=10): @@ -29,7 +29,7 @@ def random_string(length=10): return "".join(random.choices(chars, k=length)) -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime): ns = random_string() task = asyncio.create_task(server_init(runtime, ns)) diff --git a/lib/bindings/python/examples/error_handling/server.py b/lib/bindings/python/examples/error_handling/server.py index 98e8e2142df9..3a05e629c43c 100644 --- a/lib/bindings/python/examples/error_handling/server.py +++ b/lib/bindings/python/examples/error_handling/server.py @@ -17,7 +17,7 @@ import uvloop -from triton_distributed.runtime import DistributedRuntime, triton_worker +from dynemo.runtime import DistributedRuntime, dynemo_worker class RequestHandler: @@ -33,9 +33,9 @@ async def generate(self, request): yield char -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime): - await init(runtime, "triton-init") + await init(runtime, "dynemo") async def init(runtime: DistributedRuntime, ns: str): diff --git a/lib/bindings/python/examples/hello_world/client.py b/lib/bindings/python/examples/hello_world/client.py index d8fe50246126..6dff1937ee30 100644 --- a/lib/bindings/python/examples/hello_world/client.py +++ b/lib/bindings/python/examples/hello_world/client.py @@ -17,12 +17,12 @@ import uvloop -from triton_distributed.runtime import DistributedRuntime, triton_worker +from dynemo.runtime import DistributedRuntime, dynemo_worker -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime): - await init(runtime, "triton-init") + await init(runtime, "dynemo") async def init(runtime: DistributedRuntime, ns: str): diff --git a/lib/bindings/python/examples/hello_world/run.py b/lib/bindings/python/examples/hello_world/run.py index 0f53a73b959f..b73d6011c1d9 100644 --- a/lib/bindings/python/examples/hello_world/run.py +++ b/lib/bindings/python/examples/hello_world/run.py @@ -21,7 +21,7 @@ from client import init as client_init from server import init as server_init -from triton_distributed.runtime import DistributedRuntime, triton_worker +from dynemo.runtime import DistributedRuntime, dynemo_worker def random_string(length=10): @@ -29,7 +29,7 @@ def random_string(length=10): return "".join(random.choices(chars, k=length)) -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime): ns = random_string() task = asyncio.create_task(server_init(runtime, ns)) diff --git a/lib/bindings/python/examples/hello_world/server.py b/lib/bindings/python/examples/hello_world/server.py index 4c5641df899e..7b0827806b1a 100644 --- a/lib/bindings/python/examples/hello_world/server.py +++ b/lib/bindings/python/examples/hello_world/server.py @@ -17,7 +17,7 @@ import uvloop -from triton_distributed.runtime import DistributedRuntime, triton_worker +from dynemo.runtime import DistributedRuntime, dynemo_worker class RequestHandler: @@ -31,9 +31,9 @@ async def generate(self, request): yield char -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime): - await init(runtime, "triton-init") + await init(runtime, "dynemo") async def init(runtime: DistributedRuntime, ns: str): diff --git a/lib/bindings/python/examples/pipeline/backend.py b/lib/bindings/python/examples/pipeline/backend.py index 09d1f0fd1c66..05bd47065543 100644 --- a/lib/bindings/python/examples/pipeline/backend.py +++ b/lib/bindings/python/examples/pipeline/backend.py @@ -17,7 +17,7 @@ import uvloop -from triton_distributed.runtime import DistributedRuntime, triton_worker +from dynemo.runtime import DistributedRuntime, dynemo_worker uvloop.install() @@ -29,7 +29,7 @@ async def generate(self, request): yield char -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime): component = runtime.namespace("examples/pipeline").component("backend") await component.create_service() diff --git a/lib/bindings/python/examples/pipeline/frontend.py b/lib/bindings/python/examples/pipeline/frontend.py index 8d4a04356fb4..8e502367d5c9 100644 --- a/lib/bindings/python/examples/pipeline/frontend.py +++ b/lib/bindings/python/examples/pipeline/frontend.py @@ -17,7 +17,7 @@ import uvloop -from triton_distributed.runtime import DistributedRuntime, triton_worker +from dynemo.runtime import DistributedRuntime, dynemo_worker uvloop.install() @@ -32,7 +32,7 @@ async def generate(self, request): yield output.get("data") -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime): # client to the next component - in this case the middle component next = ( diff --git a/lib/bindings/python/examples/pipeline/middle.py b/lib/bindings/python/examples/pipeline/middle.py index ad206be0dfb7..9d819c38b810 100644 --- a/lib/bindings/python/examples/pipeline/middle.py +++ b/lib/bindings/python/examples/pipeline/middle.py @@ -17,7 +17,7 @@ import uvloop -from triton_distributed.runtime import DistributedRuntime, triton_worker +from dynemo.runtime import DistributedRuntime, dynemo_worker uvloop.install() @@ -32,7 +32,7 @@ async def generate(self, request): yield output.get("data") -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime): # client to backend backend = ( diff --git a/lib/bindings/python/examples/pipeline/pipeline.py b/lib/bindings/python/examples/pipeline/pipeline.py index e22ca048058c..e6626db28ed4 100644 --- a/lib/bindings/python/examples/pipeline/pipeline.py +++ b/lib/bindings/python/examples/pipeline/pipeline.py @@ -17,12 +17,12 @@ import uvloop -from triton_distributed.runtime import DistributedRuntime, triton_worker +from dynemo.runtime import DistributedRuntime, dynemo_worker uvloop.install() -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime): """ # Pipeline Example diff --git a/lib/bindings/python/examples/typed/client.py b/lib/bindings/python/examples/typed/client.py index b0e262f0b584..c65b08475d79 100644 --- a/lib/bindings/python/examples/typed/client.py +++ b/lib/bindings/python/examples/typed/client.py @@ -17,18 +17,16 @@ from protocol import Request -from triton_distributed.runtime import DistributedRuntime, triton_worker +from dynemo.runtime import DistributedRuntime, dynemo_worker -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime): """ Instantiate a `backend` client and call the `generate` endpoint """ # get endpoint - endpoint = ( - runtime.namespace("triton-init").component("backend").endpoint("generate") - ) + endpoint = runtime.namespace("dynemo").component("backend").endpoint("generate") # create client client = await endpoint.client() diff --git a/lib/bindings/python/examples/typed/server.py b/lib/bindings/python/examples/typed/server.py index 4c73f611dc0c..fe0c82b41632 100644 --- a/lib/bindings/python/examples/typed/server.py +++ b/lib/bindings/python/examples/typed/server.py @@ -19,11 +19,7 @@ import uvloop from protocol import Request, Response -from triton_distributed.runtime import ( - DistributedRuntime, - triton_endpoint, - triton_worker, -) +from dynemo.runtime import DistributedRuntime, dynemo_endpoint, dynemo_worker uvloop.install() @@ -33,19 +29,19 @@ class RequestHandler: Request handler for the generate endpoint """ - @triton_endpoint(Request, Response) + @dynemo_endpoint(Request, Response) async def generate(self, request): for char in request.data: yield char -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime): """ Instantiate a `backend` component and serve the `generate` endpoint A `Component` can serve multiple endpoints """ - component = runtime.namespace("triton-init").component("backend") + component = runtime.namespace("dynemo").component("backend") await component.create_service() endpoint = component.endpoint("generate") diff --git a/lib/bindings/python/rust/engine.rs b/lib/bindings/python/rust/engine.rs index fbc76b8aeffa..30fc019f2696 100644 --- a/lib/bindings/python/rust/engine.rs +++ b/lib/bindings/python/rust/engine.rs @@ -15,8 +15,7 @@ use std::sync::Arc; -pub use serde::{Deserialize, Serialize}; -pub use triton_distributed_runtime::{ +pub use dynemo_runtime::{ error, pipeline::{ async_trait, AsyncEngine, AsyncEngineContextProvider, Data, ManyOut, ResponseStream, @@ -25,6 +24,7 @@ pub use triton_distributed_runtime::{ protocols::annotated::Annotated, Error, Result, }; +pub use serde::{Deserialize, Serialize}; use pyo3::prelude::*; use pyo3_async_runtimes::TaskLocals; @@ -192,7 +192,6 @@ where // tell the python async generator to stop generating // right now, this is impossible as we are not passing the context to the python async generator // todo: add task-local context to the python async generator - // see: https://github.com/triton-inference-server/triton_distributed/issues/130 ctx.stop_generating(); let msg = format!("critical error: invalid response object from python async generator; application-logic-mismatch: {}", e); tracing::error!(request_id, "{}", msg); diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index d5d605a9368d..f60afc104b9e 100644 --- a/lib/bindings/python/rust/lib.rs +++ b/lib/bindings/python/rust/lib.rs @@ -25,14 +25,14 @@ use std::{fmt::Display, sync::Arc}; use tokio::sync::Mutex; use tracing_subscriber::FmtSubscriber; -use triton_distributed_runtime::{ +use dynemo_runtime::{ self as rs, pipeline::{EngineStream, ManyOut, SingleIn}, protocols::annotated::Annotated as RsAnnotated, traits::DistributedRuntimeProvider, }; -use triton_distributed_llm::{self as llm_rs}; +use dynemo_llm::{self as llm_rs}; mod engine; mod llm; diff --git a/lib/bindings/python/rust/llm/backend.rs b/lib/bindings/python/rust/llm/backend.rs index 2c5f51f84dc4..da24b3cbb46d 100644 --- a/lib/bindings/python/rust/llm/backend.rs +++ b/lib/bindings/python/rust/llm/backend.rs @@ -19,7 +19,7 @@ use crate::llm::model_card::ModelDeploymentCard; use llm_rs::protocols::common::llm_backend::{BackendInput, BackendOutput}; use llm_rs::types::Annotated; -use triton_distributed_runtime::pipeline::{Operator, ServiceBackend, ServiceFrontend, Source}; +use dynemo_runtime::pipeline::{Operator, ServiceBackend, ServiceFrontend, Source}; use crate::engine::PythonAsyncEngine; diff --git a/lib/bindings/python/rust/llm/preprocessor.rs b/lib/bindings/python/rust/llm/preprocessor.rs index e60e1e2ef665..b89740da09ff 100644 --- a/lib/bindings/python/rust/llm/preprocessor.rs +++ b/lib/bindings/python/rust/llm/preprocessor.rs @@ -27,9 +27,9 @@ use llm_rs::{ }, }; -use triton_distributed_runtime::pipeline::{Operator, ServiceFrontend, Source}; +use dynemo_runtime::pipeline::{Operator, ServiceFrontend, Source}; -use triton_distributed_runtime::pipeline::{ManyOut, SegmentSink, SingleIn}; +use dynemo_runtime::pipeline::{ManyOut, SegmentSink, SingleIn}; #[pyclass] pub(crate) struct OAIChatPreprocessor { diff --git a/lib/bindings/python/src/triton_distributed/_core.pyi b/lib/bindings/python/src/dynemo/_core.pyi similarity index 98% rename from lib/bindings/python/src/triton_distributed/_core.pyi rename to lib/bindings/python/src/dynemo/_core.pyi index 4043e590d3d7..766e380233ce 100644 --- a/lib/bindings/python/src/triton_distributed/_core.pyi +++ b/lib/bindings/python/src/dynemo/_core.pyi @@ -26,7 +26,7 @@ RequestHandler = Callable[[JsonLike], AsyncGenerator[JsonLike, None]] class DistributedRuntime: """ - The runtime object for a distributed NOVA applications + The runtime object for dynemo applications """ ... diff --git a/lib/bindings/python/src/triton_distributed/llm/__init__.py b/lib/bindings/python/src/dynemo/llm/__init__.py similarity index 83% rename from lib/bindings/python/src/triton_distributed/llm/__init__.py rename to lib/bindings/python/src/dynemo/llm/__init__.py index 87e54c415b4a..2c9abcc65e83 100644 --- a/lib/bindings/python/src/triton_distributed/llm/__init__.py +++ b/lib/bindings/python/src/dynemo/llm/__init__.py @@ -13,5 +13,5 @@ # See the License for the specific language governing permissions and # limitations under the License. -from triton_distributed._core import KvMetricsPublisher as KvMetricsPublisher -from triton_distributed._core import KvRouter as KvRouter +from dynemo._core import KvMetricsPublisher as KvMetricsPublisher +from dynemo._core import KvRouter as KvRouter diff --git a/lib/bindings/python/src/triton_distributed/runtime/__init__.py b/lib/bindings/python/src/dynemo/runtime/__init__.py similarity index 87% rename from lib/bindings/python/src/triton_distributed/runtime/__init__.py rename to lib/bindings/python/src/dynemo/runtime/__init__.py index ced5b224159f..4ae68500a4a0 100644 --- a/lib/bindings/python/src/triton_distributed/runtime/__init__.py +++ b/lib/bindings/python/src/dynemo/runtime/__init__.py @@ -22,15 +22,15 @@ # List all the classes in the _core module for re-export # import * causes "unable to detect undefined names" -from triton_distributed._core import Backend as Backend -from triton_distributed._core import Client as Client -from triton_distributed._core import DistributedRuntime as DistributedRuntime -from triton_distributed._core import KvRouter as KvRouter -from triton_distributed._core import ModelDeploymentCard as ModelDeploymentCard -from triton_distributed._core import OAIChatPreprocessor as OAIChatPreprocessor +from dynemo._core import Backend as Backend +from dynemo._core import Client as Client +from dynemo._core import DistributedRuntime as DistributedRuntime +from dynemo._core import KvRouter as KvRouter +from dynemo._core import ModelDeploymentCard as ModelDeploymentCard +from dynemo._core import OAIChatPreprocessor as OAIChatPreprocessor -def triton_worker(): +def dynemo_worker(): def decorator(func): @wraps(func) async def wrapper(*args, **kwargs): @@ -59,7 +59,7 @@ async def wrapper(*args, **kwargs): return decorator -def triton_endpoint( +def dynemo_endpoint( request_model: Union[Type[BaseModel], Type[Any]], response_model: Type[BaseModel] ) -> Callable: def decorator( diff --git a/lib/bindings/python/tests/soak.py b/lib/bindings/python/tests/soak.py index 2edd4f1ca482..8fd156f5e3ed 100644 --- a/lib/bindings/python/tests/soak.py +++ b/lib/bindings/python/tests/soak.py @@ -19,7 +19,7 @@ import uvloop -from triton_distributed.runtime import DistributedRuntime, triton_worker +from dynemo.runtime import DistributedRuntime, dynemo_worker # Soak Test # @@ -31,7 +31,7 @@ # could still eventually be a problem. -@triton_worker() +@dynemo_worker() async def worker(runtime: DistributedRuntime): ns = random_string() task = asyncio.create_task(server_init(runtime, ns)) diff --git a/lib/bindings/python/tests/test_bindings_install.py b/lib/bindings/python/tests/test_bindings_install.py index 495770cd4dc0..4056ea2bb8c7 100644 --- a/lib/bindings/python/tests/test_bindings_install.py +++ b/lib/bindings/python/tests/test_bindings_install.py @@ -20,7 +20,7 @@ def test_bindings_install(): # Verify python bindings to rust can be imported - import triton_distributed.runtime as tdr + import dynemo.runtime as tdr # Placeholder to avoid unused import errors or removal by linters assert tdr diff --git a/lib/bindings/python/tests/test_etcd_bindings.py b/lib/bindings/python/tests/test_etcd_bindings.py index f6f79ec94bc1..be4691e76ffd 100644 --- a/lib/bindings/python/tests/test_etcd_bindings.py +++ b/lib/bindings/python/tests/test_etcd_bindings.py @@ -15,7 +15,7 @@ import asyncio -from triton_distributed._core import DistributedRuntime +from dynemo._core import DistributedRuntime async def test_simple_put_get(): diff --git a/lib/llm/Cargo.lock b/lib/llm/Cargo.lock index 02bcb4ab0649..7449343a4b3f 100644 --- a/lib/llm/Cargo.lock +++ b/lib/llm/Cargo.lock @@ -1387,6 +1387,109 @@ dependencies = [ "reborrow", ] +[[package]] +name = "dynemo-llm" +version = "0.2.1" +dependencies = [ + "anyhow", + "async-openai", + "async-stream", + "async-trait", + "async_zmq", + "axum 0.8.1", + "bindgen 0.70.1", + "blake3", + "bs62", + "bytes", + "chrono", + "cmake", + "derive_builder", + "dynemo-runtime", + "either", + "erased-serde", + "futures", + "galil-seiferas", + "hf-hub 0.4.1", + "indexmap 2.7.1", + "insta", + "itertools 0.14.0", + "libc", + "llama-cpp-2", + "minijinja", + "minijinja-contrib", + "mistralrs", + "prometheus", + "proptest", + "pyo3", + "regex", + "reqwest", + "rstest", + "semver", + "sentencepiece", + "serde", + "serde-pickle", + "serde_json", + "serde_repr", + "strum 0.27.1", + "tempfile", + "thiserror 2.0.11", + "tokenizers", + "tokio", + "tokio-stream", + "tokio-util", + "toktrie 0.6.29", + "toktrie_hf_tokenizers 0.6.29", + "tracing", + "unicode-segmentation", + "uuid 1.14.0", + "validator", + "xxhash-rust", +] + +[[package]] +name = "dynemo-runtime" +version = "0.2.1" +dependencies = [ + "anyhow", + "async-nats", + "async-once-cell", + "async-stream", + "async-trait", + "async_zmq", + "blake3", + "bytes", + "chrono", + "derive-getters", + "derive_builder", + "educe", + "either", + "etcd-client", + "figment", + "futures", + "humantime", + "local-ip-address", + "log", + "nid", + "nix", + "nuid", + "once_cell", + "prometheus", + "rand", + "regex", + "serde", + "serde_json", + "socket2", + "thiserror 1.0.69", + "tokio", + "tokio-stream", + "tokio-util", + "tracing", + "tracing-subscriber", + "uuid 1.14.0", + "validator", + "xxhash-rust", +] + [[package]] name = "ed25519" version = "2.2.3" @@ -5747,109 +5850,6 @@ dependencies = [ "tracing-serde", ] -[[package]] -name = "triton-distributed-llm" -version = "0.2.1" -dependencies = [ - "anyhow", - "async-openai", - "async-stream", - "async-trait", - "async_zmq", - "axum 0.8.1", - "bindgen 0.70.1", - "blake3", - "bs62", - "bytes", - "chrono", - "cmake", - "derive_builder", - "either", - "erased-serde", - "futures", - "galil-seiferas", - "hf-hub 0.4.1", - "indexmap 2.7.1", - "insta", - "itertools 0.14.0", - "libc", - "llama-cpp-2", - "minijinja", - "minijinja-contrib", - "mistralrs", - "prometheus", - "proptest", - "pyo3", - "regex", - "reqwest", - "rstest", - "semver", - "sentencepiece", - "serde", - "serde-pickle", - "serde_json", - "serde_repr", - "strum 0.27.1", - "tempfile", - "thiserror 2.0.11", - "tokenizers", - "tokio", - "tokio-stream", - "tokio-util", - "toktrie 0.6.29", - "toktrie_hf_tokenizers 0.6.29", - "tracing", - "triton-distributed-runtime", - "unicode-segmentation", - "uuid 1.14.0", - "validator", - "xxhash-rust", -] - -[[package]] -name = "triton-distributed-runtime" -version = "0.2.1" -dependencies = [ - "anyhow", - "async-nats", - "async-once-cell", - "async-stream", - "async-trait", - "async_zmq", - "blake3", - "bytes", - "chrono", - "derive-getters", - "derive_builder", - "educe", - "either", - "etcd-client", - "figment", - "futures", - "humantime", - "local-ip-address", - "log", - "nid", - "nix", - "nuid", - "once_cell", - "prometheus", - "rand", - "regex", - "serde", - "serde_json", - "socket2", - "thiserror 1.0.69", - "tokio", - "tokio-stream", - "tokio-util", - "tracing", - "tracing-subscriber", - "uuid 1.14.0", - "validator", - "xxhash-rust", -] - [[package]] name = "try-lock" version = "0.2.5" diff --git a/lib/llm/Cargo.toml b/lib/llm/Cargo.toml index c86b73d71ed0..1a7210d99403 100644 --- a/lib/llm/Cargo.toml +++ b/lib/llm/Cargo.toml @@ -18,11 +18,11 @@ version = "0.2.1" edition = "2021" authors = ["NVIDIA"] license = "Apache-2.0" -homepage = "https://github.com/triton-inference-server/triton_distributed" -repository = "https://github.com/triton-inference-server/triton_distributed" +homepage = "https://github.com/dynemo-ai/dynemo" +repository = "https://github.com/dynemo-ai/dynemo.git" [package] -name = "triton-distributed-llm" +name = "dynemo-llm" version.workspace = true edition.workspace = true authors.workspace = true @@ -43,7 +43,7 @@ vulkan = ["llama-cpp-2/vulkan"] [workspace.dependencies] # local or crates.io -triton-distributed-runtime = { version = "0.2.0", path = "../runtime" } +dynemo-runtime = { version = "0.2.0", path = "../runtime" } # crates.io anyhow = { version = "1" } @@ -66,7 +66,7 @@ strum = { version = "0.27", features = ["derive"] } [dependencies] # repo -triton-distributed-runtime = { workspace = true } +dynemo-runtime = { workspace = true } # workspace anyhow = { workspace = true } diff --git a/lib/llm/src/backend.rs b/lib/llm/src/backend.rs index ff72656ca97c..7fb801906ea5 100644 --- a/lib/llm/src/backend.rs +++ b/lib/llm/src/backend.rs @@ -34,7 +34,7 @@ use futures::stream::{self, StreamExt}; use tracing as log; use crate::model_card::model::{ModelDeploymentCard, TokenizerKind}; -use triton_distributed_runtime::{ +use dynemo_runtime::{ pipeline::{ async_trait, AsyncEngineContextProvider, ManyOut, Operator, ResponseStream, ServerStreamingEngine, SingleIn, diff --git a/lib/llm/src/engines/llamacpp.rs b/lib/llm/src/engines/llamacpp.rs index 91646f484865..57db8c3a85ec 100644 --- a/lib/llm/src/engines/llamacpp.rs +++ b/lib/llm/src/engines/llamacpp.rs @@ -22,6 +22,11 @@ use std::{ use anyhow::Context; use async_stream::stream; use async_trait::async_trait; +use dynemo_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream}; +use dynemo_runtime::pipeline::error as pipeline_error; +use dynemo_runtime::pipeline::{Error, ManyOut, SingleIn}; +use dynemo_runtime::protocols::annotated::Annotated; +use dynemo_runtime::CancellationToken; use llama_cpp_2::{ context::{params::LlamaContextParams, LlamaContext}, llama_backend::LlamaBackend, @@ -30,11 +35,6 @@ use llama_cpp_2::{ sampling::LlamaSampler, token::LlamaToken, }; -use triton_distributed_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream}; -use triton_distributed_runtime::pipeline::error as pipeline_error; -use triton_distributed_runtime::pipeline::{Error, ManyOut, SingleIn}; -use triton_distributed_runtime::protocols::annotated::Annotated; -use triton_distributed_runtime::CancellationToken; use crate::backend::ExecutionContext; use crate::protocols::common::llm_backend::{BackendInput, LLMEngineOutput}; diff --git a/lib/llm/src/engines/mistralrs.rs b/lib/llm/src/engines/mistralrs.rs index e797ce0d9e82..3ce19777ee65 100644 --- a/lib/llm/src/engines/mistralrs.rs +++ b/lib/llm/src/engines/mistralrs.rs @@ -28,10 +28,10 @@ use mistralrs::{ }; use tokio::sync::mpsc::channel; -use triton_distributed_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream}; -use triton_distributed_runtime::pipeline::error as pipeline_error; -use triton_distributed_runtime::pipeline::{Error, ManyOut, SingleIn}; -use triton_distributed_runtime::protocols::annotated::Annotated; +use dynemo_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream}; +use dynemo_runtime::pipeline::error as pipeline_error; +use dynemo_runtime::pipeline::{Error, ManyOut, SingleIn}; +use dynemo_runtime::protocols::annotated::Annotated; use crate::protocols::openai::chat_completions::{ NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse, diff --git a/lib/llm/src/engines/sglang.rs b/lib/llm/src/engines/sglang.rs index 047a97b4a49e..89165ebde639 100644 --- a/lib/llm/src/engines/sglang.rs +++ b/lib/llm/src/engines/sglang.rs @@ -17,8 +17,8 @@ use std::path::Path; use std::sync::Arc; use crate::backend::ExecutionContext; -use triton_distributed_runtime::pipeline::error as pipeline_error; -use triton_distributed_runtime::CancellationToken; +use dynemo_runtime::pipeline::error as pipeline_error; +use dynemo_runtime::CancellationToken; mod worker; diff --git a/lib/llm/src/engines/sglang/engine.rs b/lib/llm/src/engines/sglang/engine.rs index 2ded01e6eb40..c337d63dcff4 100644 --- a/lib/llm/src/engines/sglang/engine.rs +++ b/lib/llm/src/engines/sglang/engine.rs @@ -19,10 +19,10 @@ use async_stream::stream; use async_trait::async_trait; use crate::protocols::common::llm_backend::{BackendInput, LLMEngineOutput}; -use triton_distributed_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream}; -use triton_distributed_runtime::pipeline::{Error, ManyOut, SingleIn}; -use triton_distributed_runtime::protocols::annotated::Annotated; -use triton_distributed_runtime::runtime::CancellationToken; +use dynemo_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream}; +use dynemo_runtime::pipeline::{Error, ManyOut, SingleIn}; +use dynemo_runtime::protocols::annotated::Annotated; +use dynemo_runtime::runtime::CancellationToken; use crate::engines::MultiNodeConfig; diff --git a/lib/llm/src/engines/sglang/worker.rs b/lib/llm/src/engines/sglang/worker.rs index bf9842430b46..ec3895c3c387 100644 --- a/lib/llm/src/engines/sglang/worker.rs +++ b/lib/llm/src/engines/sglang/worker.rs @@ -37,8 +37,8 @@ use tokio::sync::mpsc::Sender; use tokio::{io::AsyncBufReadExt, sync::mpsc::error::SendError}; use tokio::{io::AsyncReadExt as _, task::JoinHandle}; -use triton_distributed_runtime::protocols::annotated::Annotated; -use triton_distributed_runtime::runtime::CancellationToken; +use dynemo_runtime::protocols::annotated::Annotated; +use dynemo_runtime::runtime::CancellationToken; use crate::engines::sglang::MultiGPUConfig; use crate::engines::MultiNodeConfig; diff --git a/lib/llm/src/engines/trtllm.rs b/lib/llm/src/engines/trtllm.rs index 2b7f67954a9c..2c42fe06bd01 100644 --- a/lib/llm/src/engines/trtllm.rs +++ b/lib/llm/src/engines/trtllm.rs @@ -16,7 +16,7 @@ use std::sync::Arc; use crate::backend::ExecutionContext; -use triton_distributed_runtime::pipeline::error as pipeline_error; +use dynemo_runtime::pipeline::error as pipeline_error; pub mod executor; diff --git a/lib/llm/src/engines/trtllm/executor/engine.rs b/lib/llm/src/engines/trtllm/executor/engine.rs index 7bd8c40e937b..d02062b566e1 100644 --- a/lib/llm/src/engines/trtllm/executor/engine.rs +++ b/lib/llm/src/engines/trtllm/executor/engine.rs @@ -15,12 +15,12 @@ use anyhow::{Error, Result}; use async_trait::async_trait; +use dynemo_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream}; +use dynemo_runtime::pipeline::{ManyOut, SingleIn}; +use dynemo_runtime::protocols::annotated::Annotated; use futures::stream; use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; -use triton_distributed_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream}; -use triton_distributed_runtime::pipeline::{ManyOut, SingleIn}; -use triton_distributed_runtime::protocols::annotated::Annotated; use super::Executor; use crate::protocols::common::llm_backend::{BackendInput, LLMEngineOutput}; diff --git a/lib/llm/src/engines/vllm.rs b/lib/llm/src/engines/vllm.rs index 33292e392ef7..370129c180da 100644 --- a/lib/llm/src/engines/vllm.rs +++ b/lib/llm/src/engines/vllm.rs @@ -19,8 +19,8 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use triton_distributed_runtime::pipeline::error as pipeline_error; -use triton_distributed_runtime::CancellationToken; +use dynemo_runtime::pipeline::error as pipeline_error; +use dynemo_runtime::CancellationToken; use crate::backend::ExecutionContext; use crate::engines::MultiNodeConfig; diff --git a/lib/llm/src/engines/vllm/engine.rs b/lib/llm/src/engines/vllm/engine.rs index 1af1fc0b30cb..0dc5ee9ad269 100644 --- a/lib/llm/src/engines/vllm/engine.rs +++ b/lib/llm/src/engines/vllm/engine.rs @@ -21,10 +21,10 @@ use async_trait::async_trait; use crate::engines::vllm::worker; use crate::engines::MultiNodeConfig; use crate::protocols::common::llm_backend::{BackendInput, LLMEngineOutput}; -use triton_distributed_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream}; -use triton_distributed_runtime::pipeline::{Error, ManyOut, SingleIn}; -use triton_distributed_runtime::protocols::annotated::Annotated; -use triton_distributed_runtime::runtime::CancellationToken; +use dynemo_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream}; +use dynemo_runtime::pipeline::{Error, ManyOut, SingleIn}; +use dynemo_runtime::protocols::annotated::Annotated; +use dynemo_runtime::runtime::CancellationToken; pub struct VllmEngine { cancel_token: CancellationToken, diff --git a/lib/llm/src/engines/vllm/ray.rs b/lib/llm/src/engines/vllm/ray.rs index 466146e8fcc4..6c148cbe3cb0 100644 --- a/lib/llm/src/engines/vllm/ray.rs +++ b/lib/llm/src/engines/vllm/ray.rs @@ -24,7 +24,7 @@ use tokio::select; use tokio::time; use tracing; -use triton_distributed_runtime::CancellationToken; +use dynemo_runtime::CancellationToken; /// Default is 16 seconds, we make it a bit shorter const RAY_STOP_TIMEOUT_SECS: u32 = 10; diff --git a/lib/llm/src/engines/vllm/worker.rs b/lib/llm/src/engines/vllm/worker.rs index ed82eb904ff8..bb03c7916ad0 100644 --- a/lib/llm/src/engines/vllm/worker.rs +++ b/lib/llm/src/engines/vllm/worker.rs @@ -19,6 +19,8 @@ use std::{ }; use async_zmq::{SinkExt, StreamExt}; +use dynemo_runtime::protocols::annotated::Annotated; +use dynemo_runtime::CancellationToken; use pyo3::{ prelude::*, types::{IntoPyDict, PyBytes, PyString}, @@ -26,8 +28,6 @@ use pyo3::{ use tokio::sync::mpsc::Sender; use tokio::task::JoinHandle; use tokio::{io::AsyncBufReadExt, sync::mpsc::error::SendError}; -use triton_distributed_runtime::protocols::annotated::Annotated; -use triton_distributed_runtime::CancellationToken; use crate::engines::MultiNodeConfig; use crate::protocols::common::llm_backend::LLMEngineOutput; diff --git a/lib/llm/src/http/service/discovery.rs b/lib/llm/src/http/service/discovery.rs index 5c3e911e1887..a1e96bf3882f 100644 --- a/lib/llm/src/http/service/discovery.rs +++ b/lib/llm/src/http/service/discovery.rs @@ -18,7 +18,7 @@ use std::sync::Arc; use serde::{Deserialize, Serialize}; use tokio::sync::mpsc::Receiver; -use triton_distributed_runtime::{ +use dynemo_runtime::{ protocols::{self, annotated::Annotated}, raise, transports::etcd::{KeyValue, WatchEvent}, diff --git a/lib/llm/src/http/service/openai.rs b/lib/llm/src/http/service/openai.rs index 24e8716ca766..b7b93468c10e 100644 --- a/lib/llm/src/http/service/openai.rs +++ b/lib/llm/src/http/service/openai.rs @@ -48,7 +48,7 @@ use crate::types::{ Annotated, }; -use triton_distributed_runtime::pipeline::{AsyncEngineContext, Context}; +use dynemo_runtime::pipeline::{AsyncEngineContext, Context}; #[derive(Serialize, Deserialize)] pub(crate) struct ErrorResponse { @@ -91,7 +91,7 @@ impl ErrorResponse { ) } - /// The OAI endpoints call an [`triton_distributed_runtime::engine::AsyncEngine`] which are specialized to return + /// The OAI endpoints call an [`dynemo_runtime::engine::AsyncEngine`] which are specialized to return /// an [`anyhow::Error`]. This method will convert the [`anyhow::Error`] into an [`HttpError`]. /// If successful, it will return the [`HttpError`] as an [`ErrorResponse::internal_server_error`] /// with the details of the error. @@ -516,7 +516,7 @@ pub fn list_models_router( path: Option, ) -> (Vec, Router) { // TODO: Why do we have this endpoint? - let custom_path = path.unwrap_or("/triton/alpha/list-models".to_string()); + let custom_path = path.unwrap_or("/dynemo/alpha/list-models".to_string()); let doc_for_custom = RouteDoc::new(axum::http::Method::GET, &custom_path); // Standard OpenAI compatible list models endpoint diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 696088dbc9bc..433f6589e11c 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -14,11 +14,11 @@ // limitations under the License. use anyhow::Result; +use dynemo_runtime::{component::Component, DistributedRuntime}; use futures::stream::StreamExt; use std::{sync::Arc, time::Duration}; use tokio_util::sync::CancellationToken; use tracing; -use triton_distributed_runtime::{component::Component, DistributedRuntime}; pub mod indexer; pub mod protocols; @@ -62,7 +62,7 @@ impl KvRouter { } pub async fn new( - nats_client: triton_distributed_runtime::transports::nats::Client, + nats_client: dynemo_runtime::transports::nats::Client, service_name: String, kv_subject: String, ) -> Result> { @@ -135,7 +135,7 @@ impl KvRouter { } async fn collect_endpoints( - nats_client: triton_distributed_runtime::transports::nats::Client, + nats_client: dynemo_runtime::transports::nats::Client, service_name: String, ep_tx: tokio::sync::mpsc::Sender, cancel: CancellationToken, diff --git a/lib/llm/src/kv_router/publisher.rs b/lib/llm/src/kv_router/publisher.rs index 9963e641b68e..baa3d6e9afda 100644 --- a/lib/llm/src/kv_router/publisher.rs +++ b/lib/llm/src/kv_router/publisher.rs @@ -15,11 +15,7 @@ use crate::kv_router::{indexer::RouterEvent, protocols::*, KV_EVENT_SUBJECT}; use async_trait::async_trait; -use futures::stream; -use std::sync::Arc; -use tokio::sync::mpsc; -use tracing as log; -use triton_distributed_runtime::{ +use dynemo_runtime::{ component::Component, pipeline::{ network::Ingress, AsyncEngine, AsyncEngineContextProvider, ManyOut, ResponseStream, @@ -28,6 +24,10 @@ use triton_distributed_runtime::{ protocols::annotated::Annotated, DistributedRuntime, Error, Result, }; +use futures::stream; +use std::sync::Arc; +use tokio::sync::mpsc; +use tracing as log; pub struct KvEventPublisher { tx: mpsc::UnboundedSender, diff --git a/lib/llm/src/kv_router/worker.rs b/lib/llm/src/kv_router/worker.rs index aa607dad9a22..f886aef39e8b 100644 --- a/lib/llm/src/kv_router/worker.rs +++ b/lib/llm/src/kv_router/worker.rs @@ -19,12 +19,12 @@ pub use crate::kv_router::protocols::ForwardPassMetrics; use anyhow::Result; use derive_builder::Builder; -use triton_distributed_runtime::pipeline::network::{ +use dynemo_runtime::pipeline::network::{ ingress::push_endpoint::PushEndpoint, PushWorkHandler, }; -use triton_distributed_runtime::transports::nats::{self, ServiceExt}; +use dynemo_runtime::transports::nats::{self, ServiceExt}; use tokio::sync::watch; use tokio_util::sync::CancellationToken; diff --git a/lib/llm/src/lib.rs b/lib/llm/src/lib.rs index d142f6a99330..77de6924169b 100644 --- a/lib/llm/src/lib.rs +++ b/lib/llm/src/lib.rs @@ -13,9 +13,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! # Triton LLM +//! # Dynemo LLM //! -//! The `triton-llm` crate is a Rust library that provides a set of traits and types for building +//! The `dynemo-llm` crate is a Rust library that provides a set of traits and types for building //! distributed LLM inference solutions. pub mod backend; diff --git a/lib/llm/src/model_card/model.rs b/lib/llm/src/model_card/model.rs index 12a32bfea0fa..4049118ccea8 100644 --- a/lib/llm/src/model_card/model.rs +++ b/lib/llm/src/model_card/model.rs @@ -37,7 +37,7 @@ use std::time::Duration; use derive_builder::Builder; -use triton_distributed_runtime::slug::Slug; +use dynemo_runtime::slug::Slug; pub const BUCKET_NAME: &str = "mdc"; diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index 1cf002aec331..ebfcc149b4c8 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -35,11 +35,11 @@ use tracing; use crate::model_card::model::{ModelDeploymentCard, ModelInfo, TokenizerKind}; use crate::preprocessor::prompt::OAIChatLikeRequest; -use triton_distributed_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream}; -use triton_distributed_runtime::pipeline::{ +use dynemo_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream}; +use dynemo_runtime::pipeline::{ async_trait, AsyncEngineContext, Error, ManyOut, Operator, SingleIn, }; -use triton_distributed_runtime::protocols::annotated::{Annotated, AnnotationsProvider}; +use dynemo_runtime::protocols::annotated::{Annotated, AnnotationsProvider}; use crate::protocols::{ common::{SamplingOptionsProvider, StopConditionsProvider}, diff --git a/lib/llm/src/protocols.rs b/lib/llm/src/protocols.rs index ceeb54b6dfa3..8517415a402d 100644 --- a/lib/llm/src/protocols.rs +++ b/lib/llm/src/protocols.rs @@ -13,10 +13,10 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! # Triton LLM Protocols +//! # Dynemo LLM Protocols //! //! This module contains the protocols, i.e. messages formats, used to exchange requests and responses -//! both publicly via the HTTP API and internally between Triton components. +//! both publicly via the HTTP API and internally between Dynemo components. //! use futures::{Stream, StreamExt}; @@ -33,10 +33,10 @@ pub type DataStream = Pin + Send + Sync>>; // TODO: This is an awkward dependency that we need to address // Originally, all the Annotated/SSE Codec bits where in the LLM protocol module; however, [Annotated] -// has become the common response envelope for triton-distributed. +// has become the common response envelope for dynemo-distributed. // We may want to move the original Annotated back here and has a Infallible conversion to the the -// ResponseEnvelop in triton-distributed. -pub use triton_distributed_runtime::protocols::annotated::Annotated; +// ResponseEnvelop in dynemo-distributed. +pub use dynemo_runtime::protocols::annotated::Annotated; /// The LLM responses have multiple different fields and nests of objects to get to the actual /// text completion returned. This trait can be applied to the `choice` level objects to extract diff --git a/lib/llm/src/protocols/common/preprocessor.rs b/lib/llm/src/protocols/common/preprocessor.rs index 8c9ea6628616..a87754b3b1f4 100644 --- a/lib/llm/src/protocols/common/preprocessor.rs +++ b/lib/llm/src/protocols/common/preprocessor.rs @@ -19,7 +19,7 @@ use serde::{Deserialize, Serialize}; use super::{SamplingOptions, StopConditions}; use crate::protocols::TokenIdType; -/// [`PreprocessedRequest`] is the internal representation of an LLM request. The [`triton-llm-preprocessor`] +/// [`PreprocessedRequest`] is the internal representation of an LLM request. The [`dynemo-llm-preprocessor`] /// crate is responsible for converting request from the public APIs to this internal representation. #[derive(Serialize, Deserialize, Debug, Clone, Builder)] pub struct PreprocessedRequest { diff --git a/lib/llm/src/protocols/openai.rs b/lib/llm/src/protocols/openai.rs index 5109debf22b3..1ed02a8ccb76 100644 --- a/lib/llm/src/protocols/openai.rs +++ b/lib/llm/src/protocols/openai.rs @@ -13,7 +13,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -/// Forward openai_api_rs::v1 to triton_distributed_llm::protocols::openai::v1 pub mod chat_completions; pub mod completions; pub mod models; diff --git a/lib/llm/src/protocols/openai/chat_completions.rs b/lib/llm/src/protocols/openai/chat_completions.rs index ba86dc6f3a95..156fc714c53d 100644 --- a/lib/llm/src/protocols/openai/chat_completions.rs +++ b/lib/llm/src/protocols/openai/chat_completions.rs @@ -17,8 +17,8 @@ use super::nvext::NvExt; use super::nvext::NvExtProvider; use super::OpenAISamplingOptionsProvider; use super::OpenAIStopConditionsProvider; +use dynemo_runtime::protocols::annotated::AnnotationsProvider; use serde::{Deserialize, Serialize}; -use triton_distributed_runtime::protocols::annotated::AnnotationsProvider; use validator::Validate; mod aggregator; diff --git a/lib/llm/src/protocols/openai/completions.rs b/lib/llm/src/protocols/openai/completions.rs index 1a136ba07add..abc6198dbbf8 100644 --- a/lib/llm/src/protocols/openai/completions.rs +++ b/lib/llm/src/protocols/openai/completions.rs @@ -31,7 +31,7 @@ use super::{ CompletionUsage, ContentProvider, OpenAISamplingOptionsProvider, OpenAIStopConditionsProvider, }; -use triton_distributed_runtime::protocols::annotated::AnnotationsProvider; +use dynemo_runtime::protocols::annotated::AnnotationsProvider; #[derive(Serialize, Deserialize, Validate, Debug, Clone)] pub struct CompletionRequest { diff --git a/lib/llm/src/tokenizers/README.md b/lib/llm/src/tokenizers/README.md index 42b040e2ef10..838f3e5878ee 100644 --- a/lib/llm/src/tokenizers/README.md +++ b/lib/llm/src/tokenizers/README.md @@ -13,7 +13,7 @@ #### HuggingFace Tokenizer ```rust -use triton_distributed_llm::tokenizers::hf::HuggingFaceTokenizer; +use dynemo_llm::tokenizers::hf::HuggingFaceTokenizer; let hf_tokenizer = HuggingFaceTokenizer::from_file("tests/data/sample-models/TinyLlama_v1.1/tokenizer.json") .expect("Failed to load HuggingFace tokenizer"); @@ -22,7 +22,7 @@ let hf_tokenizer = HuggingFaceTokenizer::from_file("tests/data/sample-models/Tin ### Encoding and Decoding Text ```rust -use triton_distributed_llm::tokenizers::{HuggingFaceTokenizer, traits::{Encoder, Decoder}}; +use dynemo_llm::tokenizers::{HuggingFaceTokenizer, traits::{Encoder, Decoder}}; let tokenizer = HuggingFaceTokenizer::from_file("tests/data/sample-models/TinyLlama_v1.1/tokenizer.json") .expect("Failed to load HuggingFace tokenizer"); @@ -40,7 +40,7 @@ assert_eq!(text, decoded_text); // Using the Sequence object for encoding and decoding -use triton_distributed_llm::tokenizers::{Sequence, Tokenizer}; +use dynemo_llm::tokenizers::{Sequence, Tokenizer}; use std::sync::{Arc, RwLock}; let tokenizer = Tokenizer::from(Arc::new(tokenizer)); @@ -51,4 +51,4 @@ sequence.append_text("Your sample text here") let delta = sequence.append_token_id(1337) .expect("Failed to append token_id"); -``` \ No newline at end of file +``` diff --git a/lib/llm/src/types.rs b/lib/llm/src/types.rs index 2881b0824542..342cdd0af9c5 100644 --- a/lib/llm/src/types.rs +++ b/lib/llm/src/types.rs @@ -19,7 +19,7 @@ pub use protocols::{Annotated, TokenIdType}; pub mod openai { use super::*; - use triton_distributed_runtime::pipeline::{ServerStreamingEngine, UnaryEngine}; + use dynemo_runtime::pipeline::{ServerStreamingEngine, UnaryEngine}; pub mod completions { use super::*; diff --git a/lib/llm/tests/aggregators.rs b/lib/llm/tests/aggregators.rs index 6d16c8139f73..7b971ec460e0 100644 --- a/lib/llm/tests/aggregators.rs +++ b/lib/llm/tests/aggregators.rs @@ -13,12 +13,12 @@ // See the License for the specific language governing permissions and // limitations under the License. -use futures::StreamExt; -use triton_distributed_llm::protocols::{ +use dynemo_llm::protocols::{ codec::{create_message_stream, Message, SseCodecError}, openai::{chat_completions::NvCreateChatCompletionResponse, completions::CompletionResponse}, ContentProvider, DataStream, }; +use futures::StreamExt; const CMPL_ROOT_PATH: &str = "tests/data/replays/meta/llama-3.1-8b-instruct/completions"; const CHAT_ROOT_PATH: &str = "tests/data/replays/meta/llama-3.1-8b-instruct/chat_completions"; diff --git a/lib/llm/tests/backend.rs b/lib/llm/tests/backend.rs index 7fd264e448a6..6fe55a166de9 100644 --- a/lib/llm/tests/backend.rs +++ b/lib/llm/tests/backend.rs @@ -13,8 +13,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use triton_distributed_llm::backend::Backend; -use triton_distributed_llm::model_card::model::ModelDeploymentCard; +use dynemo_llm::backend::Backend; +use dynemo_llm::model_card::model::ModelDeploymentCard; #[tokio::test] async fn test_sequence_factory() { diff --git a/lib/llm/tests/http-service.rs b/lib/llm/tests/http-service.rs index b1c934a3886d..c58b64b994a7 100644 --- a/lib/llm/tests/http-service.rs +++ b/lib/llm/tests/http-service.rs @@ -15,28 +15,28 @@ use anyhow::Error; use async_stream::stream; -use prometheus::{proto::MetricType, Registry}; -use reqwest::StatusCode; -use std::sync::Arc; -use triton_distributed_llm::http::service::{ +use dynemo_llm::http::service::{ error::HttpError, metrics::{Endpoint, RequestType, Status}, service_v2::HttpService, Metrics, }; -use triton_distributed_llm::protocols::{ +use dynemo_llm::protocols::{ openai::{ chat_completions::{NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse}, completions::{CompletionRequest, CompletionResponse}, }, Annotated, }; -use triton_distributed_runtime::{ +use dynemo_runtime::{ pipeline::{ async_trait, AsyncEngine, AsyncEngineContextProvider, ManyOut, ResponseStream, SingleIn, }, CancellationToken, }; +use prometheus::{proto::MetricType, Registry}; +use reqwest::StatusCode; +use std::sync::Arc; struct CounterEngine {} diff --git a/lib/llm/tests/model_card.rs b/lib/llm/tests/model_card.rs index ded041fcbdae..b04afd86e34b 100644 --- a/lib/llm/tests/model_card.rs +++ b/lib/llm/tests/model_card.rs @@ -13,10 +13,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use dynemo_llm::model_card::model::{ModelDeploymentCard, PromptFormatterArtifact, TokenizerKind}; use tempfile::tempdir; -use triton_distributed_llm::model_card::model::{ - ModelDeploymentCard, PromptFormatterArtifact, TokenizerKind, -}; const HF_PATH: &str = "tests/data/sample-models/TinyLlama_v1.1"; diff --git a/lib/llm/tests/openai_completions.rs b/lib/llm/tests/openai_completions.rs index 217b702cabc3..4ba4c40ea6b5 100644 --- a/lib/llm/tests/openai_completions.rs +++ b/lib/llm/tests/openai_completions.rs @@ -14,8 +14,8 @@ // limitations under the License. use async_openai::types::CreateCompletionRequestArgs; +use dynemo_llm::protocols::openai::{self, completions::CompletionRequest}; use serde::{Deserialize, Serialize}; -use triton_distributed_llm::protocols::openai::{self, completions::CompletionRequest}; #[derive(Serialize, Deserialize, Debug, Clone)] struct CompletionSample { diff --git a/lib/llm/tests/preprocessor.rs b/lib/llm/tests/preprocessor.rs index 945076f90207..4ad814dd79d7 100644 --- a/lib/llm/tests/preprocessor.rs +++ b/lib/llm/tests/preprocessor.rs @@ -15,10 +15,10 @@ use anyhow::Ok; +use dynemo_llm::model_card::model::{ModelDeploymentCard, PromptContextMixin}; +use dynemo_llm::preprocessor::prompt::PromptFormatter; +use dynemo_llm::protocols::openai::chat_completions::NvCreateChatCompletionRequest; use serde::{Deserialize, Serialize}; -use triton_distributed_llm::model_card::model::{ModelDeploymentCard, PromptContextMixin}; -use triton_distributed_llm::preprocessor::prompt::PromptFormatter; -use triton_distributed_llm::protocols::openai::chat_completions::NvCreateChatCompletionRequest; use hf_hub::{api::tokio::ApiBuilder, Cache, Repo, RepoType}; diff --git a/lib/llm/tests/snapshots/openai_completions__minimum_viable_request.snap b/lib/llm/tests/snapshots/openai_completions__minimum_viable_request.snap index cf8dc2485e3c..b7e6e7f477c3 100644 --- a/lib/llm/tests/snapshots/openai_completions__minimum_viable_request.snap +++ b/lib/llm/tests/snapshots/openai_completions__minimum_viable_request.snap @@ -1,5 +1,5 @@ --- -source: triton-llm/tests/openai_completions.rs +source: dynemo-llm/tests/openai_completions.rs expression: request --- { diff --git a/lib/llm/tests/snapshots/openai_completions__valid_samples-10.snap b/lib/llm/tests/snapshots/openai_completions__valid_samples-10.snap index c644584b6431..67a6c508881d 100644 --- a/lib/llm/tests/snapshots/openai_completions__valid_samples-10.snap +++ b/lib/llm/tests/snapshots/openai_completions__valid_samples-10.snap @@ -1,5 +1,5 @@ --- -source: triton-llm/tests/openai_completions.rs +source: dynemo-llm/tests/openai_completions.rs description: "should have prompt, model, and logit_bias fields with the logits_bias having two key/value pairs" expression: sample.request --- diff --git a/lib/llm/tests/snapshots/openai_completions__valid_samples-3.snap b/lib/llm/tests/snapshots/openai_completions__valid_samples-3.snap index 650d0e0e1700..f3572a755bfd 100644 --- a/lib/llm/tests/snapshots/openai_completions__valid_samples-3.snap +++ b/lib/llm/tests/snapshots/openai_completions__valid_samples-3.snap @@ -1,5 +1,5 @@ --- -source: triton-llm/tests/openai_completions.rs +source: dynemo-llm/tests/openai_completions.rs description: "should have prompt, model, and temperature fields" expression: sample.request --- diff --git a/lib/llm/tests/snapshots/openai_completions__valid_samples-4.snap b/lib/llm/tests/snapshots/openai_completions__valid_samples-4.snap index 91060ef76dc1..6b14a57a3e6c 100644 --- a/lib/llm/tests/snapshots/openai_completions__valid_samples-4.snap +++ b/lib/llm/tests/snapshots/openai_completions__valid_samples-4.snap @@ -1,5 +1,5 @@ --- -source: triton-llm/tests/openai_completions.rs +source: dynemo-llm/tests/openai_completions.rs description: "should have prompt, model, and top_p fields" expression: sample.request --- diff --git a/lib/llm/tests/snapshots/openai_completions__valid_samples-7.snap b/lib/llm/tests/snapshots/openai_completions__valid_samples-7.snap index 8fd982813e2c..40fba7575e08 100644 --- a/lib/llm/tests/snapshots/openai_completions__valid_samples-7.snap +++ b/lib/llm/tests/snapshots/openai_completions__valid_samples-7.snap @@ -1,5 +1,5 @@ --- -source: triton-llm/tests/openai_completions.rs +source: dynemo-llm/tests/openai_completions.rs description: "should have prompt, model, and stop fields" expression: sample.request --- diff --git a/lib/llm/tests/snapshots/openai_completions__valid_samples-9.snap b/lib/llm/tests/snapshots/openai_completions__valid_samples-9.snap index 2e828b70ea5a..169b9948d01d 100644 --- a/lib/llm/tests/snapshots/openai_completions__valid_samples-9.snap +++ b/lib/llm/tests/snapshots/openai_completions__valid_samples-9.snap @@ -1,5 +1,5 @@ --- -source: triton-llm/tests/openai_completions.rs +source: dynemo-llm/tests/openai_completions.rs description: "should have prompt, model, and stream fields" expression: sample.request --- diff --git a/lib/llm/tests/snapshots/openai_completions__valid_samples.snap b/lib/llm/tests/snapshots/openai_completions__valid_samples.snap index db4534d8ea5e..e816aa249316 100644 --- a/lib/llm/tests/snapshots/openai_completions__valid_samples.snap +++ b/lib/llm/tests/snapshots/openai_completions__valid_samples.snap @@ -1,5 +1,5 @@ --- -source: triton-llm/tests/openai_completions.rs +source: dynemo-llm/tests/openai_completions.rs description: should have only prompt and model fields expression: sample.request --- diff --git a/lib/llm/tests/snapshots/preprocessor__mulit_turn_with_continuation@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap b/lib/llm/tests/snapshots/preprocessor__mulit_turn_with_continuation@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap index 3c64470ddc88..2ce1dea08327 100644 --- a/lib/llm/tests/snapshots/preprocessor__mulit_turn_with_continuation@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap +++ b/lib/llm/tests/snapshots/preprocessor__mulit_turn_with_continuation@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap @@ -1,5 +1,5 @@ --- -source: triton-llm/tests/preprocessor.rs +source: dynemo-llm/tests/preprocessor.rs expression: formatted_prompt info: messages: diff --git a/lib/llm/tests/snapshots/preprocessor__mulit_turn_with_system@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap b/lib/llm/tests/snapshots/preprocessor__mulit_turn_with_system@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap index 3f70453d65d6..64b24b0ee707 100644 --- a/lib/llm/tests/snapshots/preprocessor__mulit_turn_with_system@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap +++ b/lib/llm/tests/snapshots/preprocessor__mulit_turn_with_system@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap @@ -1,5 +1,5 @@ --- -source: triton-llm/tests/preprocessor.rs +source: dynemo-llm/tests/preprocessor.rs expression: formatted_prompt info: messages: diff --git a/lib/llm/tests/snapshots/preprocessor__mulit_turn_with_system_with_tools@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap b/lib/llm/tests/snapshots/preprocessor__mulit_turn_with_system_with_tools@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap index f06a197d60f0..e60f8e470188 100644 --- a/lib/llm/tests/snapshots/preprocessor__mulit_turn_with_system_with_tools@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap +++ b/lib/llm/tests/snapshots/preprocessor__mulit_turn_with_system_with_tools@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap @@ -1,5 +1,5 @@ --- -source: triton-llm/tests/preprocessor.rs +source: dynemo-llm/tests/preprocessor.rs expression: formatted_prompt info: messages: diff --git a/lib/llm/tests/snapshots/preprocessor__mulit_turn_without_system@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap b/lib/llm/tests/snapshots/preprocessor__mulit_turn_without_system@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap index dec76dc71a94..953c75fb8b27 100644 --- a/lib/llm/tests/snapshots/preprocessor__mulit_turn_without_system@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap +++ b/lib/llm/tests/snapshots/preprocessor__mulit_turn_without_system@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap @@ -1,5 +1,5 @@ --- -source: triton-llm/tests/preprocessor.rs +source: dynemo-llm/tests/preprocessor.rs expression: formatted_prompt info: messages: diff --git a/lib/llm/tests/snapshots/preprocessor__multi_turn_with_continuation@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap b/lib/llm/tests/snapshots/preprocessor__multi_turn_with_continuation@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap index 3c64470ddc88..2ce1dea08327 100644 --- a/lib/llm/tests/snapshots/preprocessor__multi_turn_with_continuation@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap +++ b/lib/llm/tests/snapshots/preprocessor__multi_turn_with_continuation@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap @@ -1,5 +1,5 @@ --- -source: triton-llm/tests/preprocessor.rs +source: dynemo-llm/tests/preprocessor.rs expression: formatted_prompt info: messages: diff --git a/lib/llm/tests/snapshots/preprocessor__multi_turn_with_system_with_tools@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap b/lib/llm/tests/snapshots/preprocessor__multi_turn_with_system_with_tools@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap index f06a197d60f0..e60f8e470188 100644 --- a/lib/llm/tests/snapshots/preprocessor__multi_turn_with_system_with_tools@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap +++ b/lib/llm/tests/snapshots/preprocessor__multi_turn_with_system_with_tools@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap @@ -1,5 +1,5 @@ --- -source: triton-llm/tests/preprocessor.rs +source: dynemo-llm/tests/preprocessor.rs expression: formatted_prompt info: messages: diff --git a/lib/llm/tests/snapshots/preprocessor__single_turn@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap b/lib/llm/tests/snapshots/preprocessor__single_turn@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap index 63f782352749..2d5246809292 100644 --- a/lib/llm/tests/snapshots/preprocessor__single_turn@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap +++ b/lib/llm/tests/snapshots/preprocessor__single_turn@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap @@ -1,5 +1,5 @@ --- -source: triton-llm/tests/preprocessor.rs +source: dynemo-llm/tests/preprocessor.rs expression: formatted_prompt info: messages: diff --git a/lib/llm/tests/snapshots/preprocessor__single_turn_with_tools@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap b/lib/llm/tests/snapshots/preprocessor__single_turn_with_tools@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap index 5c243b910565..d21e2f3f3563 100644 --- a/lib/llm/tests/snapshots/preprocessor__single_turn_with_tools@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap +++ b/lib/llm/tests/snapshots/preprocessor__single_turn_with_tools@meta_llama_llama_3_1_70b_instruct__1605565_e45e5991.snap @@ -1,5 +1,5 @@ --- -source: triton-llm/tests/preprocessor.rs +source: dynemo-llm/tests/preprocessor.rs expression: formatted_prompt info: messages: diff --git a/lib/llm/tests/tokenizers.rs b/lib/llm/tests/tokenizers.rs index f7d107f64169..6fbc7af9c52f 100644 --- a/lib/llm/tests/tokenizers.rs +++ b/lib/llm/tests/tokenizers.rs @@ -25,10 +25,10 @@ //! in a hashmap. We will then use these hashes to test that the tokenizer is working correctly. This //! will detect if upstream dependency changes result in different/new behavior. +use dynemo_llm::tokenizers::traits::{Decoder, Encoder}; +use dynemo_llm::tokenizers::*; use std::collections::HashMap; use std::sync::Arc; -use triton_distributed_llm::tokenizers::traits::{Decoder, Encoder}; -use triton_distributed_llm::tokenizers::*; const TEST_PROMPTS: [&str; 4] = [ "deep learning is", diff --git a/lib/runtime/Cargo.lock b/lib/runtime/Cargo.lock index 3e151536a069..577ccadbff87 100644 --- a/lib/runtime/Cargo.lock +++ b/lib/runtime/Cargo.lock @@ -695,6 +695,54 @@ dependencies = [ "syn 2.0.98", ] +[[package]] +name = "dynemo-runtime" +version = "0.2.1" +dependencies = [ + "anyhow", + "assert_matches", + "async-nats", + "async-once-cell", + "async-stream", + "async-trait", + "async_zmq", + "blake3", + "bytes", + "chrono", + "derive-getters", + "derive_builder", + "educe", + "either", + "env_logger", + "etcd-client", + "figment", + "futures", + "humantime", + "local-ip-address", + "log", + "nid", + "nix", + "nuid", + "once_cell", + "prometheus", + "rand", + "regex", + "rstest", + "serde", + "serde_json", + "socket2", + "temp-env", + "thiserror", + "tokio", + "tokio-stream", + "tokio-util", + "tracing", + "tracing-subscriber", + "uuid", + "validator", + "xxhash-rust", +] + [[package]] name = "ed25519" version = "2.2.3" @@ -2930,54 +2978,6 @@ dependencies = [ "tracing-serde", ] -[[package]] -name = "triton-distributed-runtime" -version = "0.2.1" -dependencies = [ - "anyhow", - "assert_matches", - "async-nats", - "async-once-cell", - "async-stream", - "async-trait", - "async_zmq", - "blake3", - "bytes", - "chrono", - "derive-getters", - "derive_builder", - "educe", - "either", - "env_logger", - "etcd-client", - "figment", - "futures", - "humantime", - "local-ip-address", - "log", - "nid", - "nix", - "nuid", - "once_cell", - "prometheus", - "rand", - "regex", - "rstest", - "serde", - "serde_json", - "socket2", - "temp-env", - "thiserror", - "tokio", - "tokio-stream", - "tokio-util", - "tracing", - "tracing-subscriber", - "uuid", - "validator", - "xxhash-rust", -] - [[package]] name = "try-lock" version = "0.2.5" diff --git a/lib/runtime/Cargo.toml b/lib/runtime/Cargo.toml index c1069091fc6e..d7f5431dc1f5 100644 --- a/lib/runtime/Cargo.toml +++ b/lib/runtime/Cargo.toml @@ -14,16 +14,16 @@ # limitations under the License. [package] -name = "triton-distributed-runtime" +name = "dynemo-runtime" description = "Distributed Inference Framework" readme = "README.md" version = "0.2.1" # TODO: Centralize Version Automation edition = "2021" authors = ["NVIDIA"] license = "Apache-2.0" -homepage = "https://github.com/triton-inference-server/triton_distributed" -repository = "https://github.com/triton-inference-server/triton_distributed" -keywords = ["llm", "genai", "inference", "nvidia", "distributed", "triton"] +homepage = "https://github.com/dynemo-ai/dynemo" +repository = "https://github.com/dynemo-ai/dynemo.git" +keywords = ["llm", "genai", "inference", "nvidia", "distributed", "dynemo"] [features] default = [] diff --git a/lib/runtime/README.md b/lib/runtime/README.md index 67a79f802dc0..34cd78aeaee2 100644 --- a/lib/runtime/README.md +++ b/lib/runtime/README.md @@ -70,8 +70,7 @@ above, you can manually launch each: When developing or running examples, any process or user that shared your core-services (`etcd` and `nats.io`) will be operating within your distributed runtime. -The current examples use a hard-coded `namespace`. We will address the `namespace` collisions in this -[issue](https://github.com/triton-inference-server/triton_distributed/issues/114). +The current examples use a hard-coded `namespace`. We will address the `namespace` collisions later. All examples require the `etcd` and `nats.io` pre-requisites to be running and available. diff --git a/lib/runtime/src/config.rs b/lib/runtime/src/config.rs index 1dbee8179bf3..56b5973e0d6f 100644 --- a/lib/runtime/src/config.rs +++ b/lib/runtime/src/config.rs @@ -35,7 +35,7 @@ impl WorkerConfig { // All calls should be global and thread safe. Figment::new() .merge(Serialized::defaults(Self::default())) - .merge(Env::prefixed("TRD_WORKER_")) + .merge(Env::prefixed("DYN_WORKER_")) .extract() .unwrap() // safety: Called on startup, so panic is reasonable } @@ -81,10 +81,10 @@ impl RuntimeConfig { pub(crate) fn figment() -> Figment { Figment::new() .merge(Serialized::defaults(RuntimeConfig::default())) - .merge(Toml::file("/opt/triton/defaults/runtime.toml")) - .merge(Toml::file("/opt/triton/etc/runtime.toml")) - .merge(Env::prefixed("TRD_RUNTIME_").filter_map(|k| { - let full_key = format!("TRD_RUNTIME_{}", k.as_str()); + .merge(Toml::file("/opt/dynemo/defaults/runtime.toml")) + .merge(Toml::file("/opt/dynemo/etc/runtime.toml")) + .merge(Env::prefixed("DYN_RUNTIME_").filter_map(|k| { + let full_key = format!("DYN_RUNTIME_{}", k.as_str()); // filters out empty environment variables match std::env::var(&full_key) { Ok(v) if !v.is_empty() => Some(k.into()), @@ -97,10 +97,10 @@ impl RuntimeConfig { /// Configuration is priorities in the following order, where the last has the lowest priority: /// 1. Environment variables (top priority) /// TO DO: Add documentation for configuration files. Paths should be configurable. - /// 2. /opt/triton/etc/runtime.toml - /// 3. /opt/triton/defaults/runtime.toml (lowest priority) + /// 2. /opt/dynemo/etc/runtime.toml + /// 3. /opt/dynemo/defaults/runtime.toml (lowest priority) /// - /// Environment variables are prefixed with `TRD_RUNTIME_` + /// Environment variables are prefixed with `DYN_RUNTIME_` pub fn from_settings() -> Result { let config: RuntimeConfig = Self::figment().extract()?; config.validate()?; @@ -159,15 +159,15 @@ pub fn is_truthy(val: &str) -> bool { } /// Check whether JSONL logging enabled -/// Set the `TRD_LOGGING_JSONL` environment variable a [`is_truthy`] value +/// Set the `DYN_LOGGING_JSONL` environment variable a [`is_truthy`] value pub fn jsonl_logging_enabled() -> bool { - env_is_truthy("TRD_LOGGING_JSONL") + env_is_truthy("DYN_LOGGING_JSONL") } /// Check whether logging with ANSI terminal escape codes and colors is disabled. -/// Set the `TRD_SDK_DISABLE_ANSI_LOGGING` environment variable a [`is_truthy`] value +/// Set the `DYN_SDK_DISABLE_ANSI_LOGGING` environment variable a [`is_truthy`] value pub fn disable_ansi_logging() -> bool { - env_is_truthy("TRD_SDK_DISABLE_ANSI_LOGGING") + env_is_truthy("DYN_SDK_DISABLE_ANSI_LOGGING") } #[cfg(test)] @@ -178,8 +178,8 @@ mod tests { fn test_runtime_config_with_env_vars() -> Result<()> { temp_env::with_vars( vec![ - ("TRD_RUNTIME_NUM_WORKER_THREADS", Some("24")), - ("TRD_RUNTIME_MAX_BLOCKING_THREADS", Some("32")), + ("DYN_RUNTIME_NUM_WORKER_THREADS", Some("24")), + ("DYN_RUNTIME_MAX_BLOCKING_THREADS", Some("32")), ], || { let config = RuntimeConfig::from_settings()?; @@ -194,8 +194,8 @@ mod tests { fn test_runtime_config_defaults() -> Result<()> { temp_env::with_vars( vec![ - ("TRD_RUNTIME_NUM_WORKER_THREADS", None::<&str>), - ("TRD_RUNTIME_MAX_BLOCKING_THREADS", Some("")), + ("DYN_RUNTIME_NUM_WORKER_THREADS", None::<&str>), + ("DYN_RUNTIME_MAX_BLOCKING_THREADS", Some("")), ], || { let config = RuntimeConfig::from_settings()?; @@ -215,8 +215,8 @@ mod tests { fn test_runtime_config_rejects_invalid_thread_count() -> Result<()> { temp_env::with_vars( vec![ - ("TRD_RUNTIME_NUM_WORKER_THREADS", Some("0")), - ("TRD_RUNTIME_MAX_BLOCKING_THREADS", Some("0")), + ("DYN_RUNTIME_NUM_WORKER_THREADS", Some("0")), + ("DYN_RUNTIME_MAX_BLOCKING_THREADS", Some("0")), ], || { let result = RuntimeConfig::from_settings(); diff --git a/lib/runtime/src/logging.rs b/lib/runtime/src/logging.rs index cd239a7ed280..9bc6bb6b7232 100644 --- a/lib/runtime/src/logging.rs +++ b/lib/runtime/src/logging.rs @@ -13,17 +13,17 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Triton Distributed Logging Module. +//! Dynemo Distributed Logging Module. //! //! - Configuration loaded from: //! 1. Environment variables (highest priority). -//! 2. Optional TOML file pointed to by the `TRD_LOGGING_CONFIG_PATH` environment variable. -//! 3. `/opt/triton/etc/logging.toml`. +//! 2. Optional TOML file pointed to by the `DYN_LOGGING_CONFIG_PATH` environment variable. +//! 3. `/opt/dynemo/etc/logging.toml`. //! //! Logging can take two forms: `READABLE` or `JSONL`. The default is `READABLE`. `JSONL` -//! can be enabled by setting the `TRD_LOGGING_JSONL` environment variable to `1`. +//! can be enabled by setting the `DYN_LOGGING_JSONL` environment variable to `1`. //! -//! Filters can be configured using the `TRD_LOG` environment variable or by setting the `filters` +//! Filters can be configured using the `DYN_LOG` environment variable or by setting the `filters` //! key in the TOML configuration file. Filters are comma-separated key-value pairs where the key //! is the crate or module name and the value is the log level. The default log level is `error`. //! @@ -53,13 +53,13 @@ use tracing_subscriber::EnvFilter; use tracing_subscriber::{filter::Directive, fmt}; /// ENV used to set the log level -const FILTER_ENV: &str = "TRD_LOG"; +const FILTER_ENV: &str = "DYN_LOG"; /// Default log level const DEFAULT_FILTER_LEVEL: &str = "error"; /// ENV used to set the path to the logging configuration file -const CONFIG_PATH_ENV: &str = "TRD_LOGGING_CONFIG_PATH"; +const CONFIG_PATH_ENV: &str = "DYN_LOGGING_CONFIG_PATH"; /// Once instance to ensure the logger is only initialized once static INIT: Once = Once::new(); @@ -155,7 +155,7 @@ fn load_config() -> LoggingConfig { let config_path = std::env::var(CONFIG_PATH_ENV).unwrap_or_else(|_| "".to_string()); let figment = Figment::new() .merge(Serialized::defaults(LoggingConfig::default())) - .merge(Toml::file("/opt/triton/etc/logging.toml")) + .merge(Toml::file("/opt/dynemo/etc/logging.toml")) .merge(Toml::file(config_path)); figment.extract().unwrap() diff --git a/lib/runtime/src/pipeline/context.rs b/lib/runtime/src/pipeline/context.rs index de5555148957..46eb58360b73 100644 --- a/lib/runtime/src/pipeline/context.rs +++ b/lib/runtime/src/pipeline/context.rs @@ -299,7 +299,7 @@ impl From> for StreamContext { } } -// TODO - refactor here - this came from the triton-llm-async-engine crate +// TODO - refactor here - this came from the dynemo-llm-async-engine crate use tokio::sync::watch::{channel, Receiver, Sender}; diff --git a/lib/runtime/src/pipeline/registry.rs b/lib/runtime/src/pipeline/registry.rs index 76abcbbd35c1..735e5bd92346 100644 --- a/lib/runtime/src/pipeline/registry.rs +++ b/lib/runtime/src/pipeline/registry.rs @@ -22,7 +22,7 @@ use std::sync::Arc; /// # Examples /// /// ``` -/// use triton_distributed_runtime::pipeline::registry::Registry; +/// use dynemo_runtime::pipeline::registry::Registry; /// /// let mut registry = Registry::new(); /// diff --git a/lib/runtime/src/protocols.rs b/lib/runtime/src/protocols.rs index 9aa192b8bc9c..0b5e5e4a1d73 100644 --- a/lib/runtime/src/protocols.rs +++ b/lib/runtime/src/protocols.rs @@ -95,7 +95,7 @@ impl From<&str> for Endpoint { /// /// # Examples /// ```ignore - /// use triton_distributed::protocols::Endpoint; + /// use dynemo_runtime:protocols::Endpoint; /// /// let endpoint = Endpoint::from("namespace/component/endpoint"); /// assert_eq!(endpoint.namespace, "namespace"); @@ -150,7 +150,7 @@ impl FromStr for Endpoint { /// # Examples /// ```ignore /// use std::str::FromStr; - /// use triton_distributed::protocols::Endpoint; + /// use dynemo_runtime:protocols::Endpoint; /// /// let endpoint: Endpoint = "namespace/component/endpoint".parse().unwrap(); /// assert_eq!(endpoint.namespace, "namespace"); diff --git a/lib/runtime/src/worker.rs b/lib/runtime/src/worker.rs index 9c967c407b03..61aff5012cd8 100644 --- a/lib/runtime/src/worker.rs +++ b/lib/runtime/src/worker.rs @@ -18,17 +18,17 @@ //! //! In the future, the [Worker] should probably be moved to a procedural macro similar //! to the `#[tokio::main]` attribute, where we might annotate an async main function with -//! `#[triton::main]` or similar. +//! `#[dynemo::main]` or similar. //! //! The [Worker::execute] method is designed to be called once from main and will block //! the calling thread until the application completes or is canceled. The method initialized //! the signal handler used to trap `SIGINT` and `SIGTERM` signals and trigger a graceful shutdown. //! //! On termination, the user application is given a graceful shutdown period of controlled by -//! the [TRD_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT] environment variable. If the application does not +//! the [DYN_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT] environment variable. If the application does not //! shutdown in time, the worker will terminate the application with an exit code of 911. //! -//! The default values of [TRD_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT] differ between the development +//! The default values of [DYN_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT] differ between the development //! and release builds. In development, the default is [DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_DEBUG] and //! in release, the default is [DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_RELEASE]. @@ -45,10 +45,10 @@ static INIT: OnceCell>>>> = Once const SHUTDOWN_MESSAGE: &str = "Application received shutdown signal; attempting to gracefully shutdown"; const SHUTDOWN_TIMEOUT_MESSAGE: &str = - "Use TRD_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT to control the graceful shutdown timeout"; + "Use DYN_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT to control the graceful shutdown timeout"; /// Environment variable to control the graceful shutdown timeout -pub const TRD_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT: &str = "TRD_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT"; +pub const DYN_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT: &str = "DYN_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT"; /// Default graceful shutdown timeout in seconds in debug mode pub const DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_DEBUG: u64 = 5; @@ -106,7 +106,7 @@ impl Worker { let primary = runtime.primary(); let secondary = runtime.secondary(); - let timeout = std::env::var(TRD_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT) + let timeout = std::env::var(DYN_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT) .ok() .and_then(|s| s.parse::().ok()) .unwrap_or({ diff --git a/lib/runtime/tests/common/engines.rs b/lib/runtime/tests/common/engines.rs index facb2c2cc0fe..815aba3b056a 100644 --- a/lib/runtime/tests/common/engines.rs +++ b/lib/runtime/tests/common/engines.rs @@ -21,12 +21,12 @@ use async_trait::async_trait; use futures::Stream; use tokio::sync::mpsc; -use triton_distributed_runtime::engine::{ +use dynemo_runtime::engine::{ AsyncEngine, AsyncEngineContext, AsyncEngineContextProvider, AsyncEngineStream, Data as DataType, Engine, EngineStream, }; -use triton_distributed_runtime::pipeline::{ +use dynemo_runtime::pipeline::{ context::{Context, StreamContext}, Error, ManyOut, SingleIn, }; diff --git a/lib/runtime/tests/common/mock.rs b/lib/runtime/tests/common/mock.rs index 740ad8f2ad57..518693a6720b 100644 --- a/lib/runtime/tests/common/mock.rs +++ b/lib/runtime/tests/common/mock.rs @@ -21,8 +21,8 @@ use futures::StreamExt; use serde::{Deserialize, Serialize}; use tokio::sync::mpsc; -use triton_distributed_runtime::engine::{AsyncEngine, AsyncEngineContext, Data, ResponseStream}; -use triton_distributed_runtime::pipeline::{ +use dynemo_runtime::engine::{AsyncEngine, AsyncEngineContext, Data, ResponseStream}; +use dynemo_runtime::pipeline::{ context::{Context, StreamContext}, Error, ManyOut, PipelineError, PipelineIO, SegmentSource, SingleIn, }; diff --git a/lib/runtime/tests/lifecycle.rs b/lib/runtime/tests/lifecycle.rs index 183c5fdbe996..ea50e7d8cee5 100644 --- a/lib/runtime/tests/lifecycle.rs +++ b/lib/runtime/tests/lifecycle.rs @@ -13,7 +13,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use triton_distributed_runtime::{worker::Worker, Result, Runtime}; +use dynemo_runtime::{worker::Worker, Result, Runtime}; async fn hello_world(_runtime: Runtime) -> Result<()> { Ok(()) diff --git a/lib/runtime/tests/pipeline.rs b/lib/runtime/tests/pipeline.rs index 1d554fa46d04..9e7b53c13fbf 100644 --- a/lib/runtime/tests/pipeline.rs +++ b/lib/runtime/tests/pipeline.rs @@ -17,8 +17,8 @@ use futures::{stream, StreamExt}; use serde::{Deserialize, Serialize}; use std::{sync::Arc, time::Duration}; -use triton_distributed_runtime::engine::ResponseStream; -use triton_distributed_runtime::{ +use dynemo_runtime::engine::ResponseStream; +use dynemo_runtime::{ pipeline::{ async_trait, AsyncEngine, Data, Event, ManyOut, Operator, ServiceBackend, ServiceEngine, ServiceFrontend, SingleIn, *, diff --git a/lib/runtime/tests/soak.rs b/lib/runtime/tests/soak.rs index 59796a3f2522..854f99152147 100644 --- a/lib/runtime/tests/soak.rs +++ b/lib/runtime/tests/soak.rs @@ -16,12 +16,9 @@ #[cfg(feature = "integration")] mod integration { - pub const DEFAULT_NAMESPACE: &str = "triton-init"; + pub const DEFAULT_NAMESPACE: &str = "dynemo"; - use futures::StreamExt; - use std::{sync::Arc, time::Duration}; - use tokio::time::Instant; - use triton_distributed_runtime::{ + use dynemo_runtime::{ logging, pipeline::{ async_trait, network::Ingress, AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, @@ -30,6 +27,9 @@ mod integration { protocols::annotated::Annotated, DistributedRuntime, ErrorContext, Result, Runtime, Worker, }; + use futures::StreamExt; + use std::{sync::Arc, time::Duration}; + use tokio::time::Instant; #[test] fn main() -> Result<()> { @@ -100,11 +100,11 @@ mod integration { async fn client(runtime: DistributedRuntime) -> Result<()> { // get the run duration from env - let run_duration = std::env::var("TRD_SOAK_RUN_DURATION").unwrap_or("1m".to_string()); + let run_duration = std::env::var("DYN_SOAK_RUN_DURATION").unwrap_or("1m".to_string()); let run_duration = humantime::parse_duration(&run_duration).unwrap_or(Duration::from_secs(60)); - let batch_load = std::env::var("TRD_SOAK_BATCH_LOAD").unwrap_or("10000".to_string()); + let batch_load = std::env::var("DYN_SOAK_BATCH_LOAD").unwrap_or("10000".to_string()); let batch_load: usize = batch_load.parse().unwrap_or(10000); let client = runtime diff --git a/pyproject.toml b/pyproject.toml index 5ed8c9d8dc91..8f7d59edf4a9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,12 +14,12 @@ # limitations under the License. [project] -name = "triton-distributed" +name = "dynemo" version = "0.2.1" description = "Distributed Inference Framework" readme = "README.md" authors = [ - { name = "NVIDIA Inc.", email = "sw-dl-triton@nvidia.com" }, + { name = "NVIDIA Inc.", email = "sw-dl-dynemo@nvidia.com" }, ] license = { file = "LICENSE" } requires-python = ">=3.10" @@ -30,11 +30,11 @@ dependencies = [ ] [tool.maturin] -module-name = "triton_distributed._core" +module-name = "dynemo._core" # Point to the Rust crate directory manifest-path = "lib/bindings/python/Cargo.toml" # Include all Python packages -python-packages = ["triton_distributed"] +python-packages = ["dynemo"] # Specify the Python source directory python-source = "lib/bindings/python/src" @@ -75,7 +75,7 @@ line_length = 88 balanced_wrapping = true indent = " " skip = ["build"] -known_first_party = ["triton_distributed"] +known_first_party = ["dynemo"] [tool.pytest.ini_options] minversion = "8.0"