llama-stack 0.0.42__py3-none-any.whl → 0.3.4__py3-none-any.whl

llama_stack/providers/inline/inference/meta_reference/inference.py

@@ -0,0 +1,158 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import asyncio
+from collections.abc import AsyncIterator
+
+from llama_stack.apis.inference import (
+    InferenceProvider,
+    OpenAIChatCompletionRequestWithExtraBody,
+    OpenAICompletionRequestWithExtraBody,
+)
+from llama_stack.apis.inference.inference import (
+    OpenAIChatCompletion,
+    OpenAIChatCompletionChunk,
+    OpenAICompletion,
+)
+from llama_stack.apis.models import Model, ModelType
+from llama_stack.log import get_logger
+from llama_stack.models.llama.llama3.chat_format import ChatFormat as Llama3ChatFormat
+from llama_stack.models.llama.llama3.tokenizer import Tokenizer as Llama3Tokenizer
+from llama_stack.models.llama.llama4.chat_format import ChatFormat as Llama4ChatFormat
+from llama_stack.models.llama.llama4.tokenizer import Tokenizer as Llama4Tokenizer
+from llama_stack.models.llama.sku_list import resolve_model
+from llama_stack.models.llama.sku_types import ModelFamily
+from llama_stack.providers.datatypes import ModelsProtocolPrivate
+from llama_stack.providers.utils.inference.embedding_mixin import (
+    SentenceTransformerEmbeddingMixin,
+)
+from llama_stack.providers.utils.inference.model_registry import (
+    ModelRegistryHelper,
+    build_hf_repo_model_entry,
+)
+
+from .config import MetaReferenceInferenceConfig
+from .generators import LlamaGenerator
+from .model_parallel import LlamaModelParallelGenerator
+
+log = get_logger(__name__, category="inference")
+# there's a single model parallel process running serving the model. for now,
+# we don't support multiple concurrent requests to this process.
+SEMAPHORE = asyncio.Semaphore(1)
+
+
+def llama_builder_fn(config: MetaReferenceInferenceConfig, model_id: str, llama_model: Model) -> LlamaGenerator:
+    return LlamaGenerator(config, model_id, llama_model)
+
+
+class MetaReferenceInferenceImpl(
+    SentenceTransformerEmbeddingMixin,
+    InferenceProvider,
+    ModelsProtocolPrivate,
+):
+    def __init__(self, config: MetaReferenceInferenceConfig) -> None:
+        self.config = config
+        self.model_id = None
+        self.llama_model = None
+
+    async def initialize(self) -> None:
+        pass
+
+    async def shutdown(self) -> None:
+        if self.config.create_distributed_process_group:
+            self.generator.stop()
+
+    async def openai_completion(
+        self,
+        params: OpenAICompletionRequestWithExtraBody,
+    ) -> OpenAICompletion:
+        raise NotImplementedError("OpenAI completion not supported by meta reference provider")
+
+    async def should_refresh_models(self) -> bool:
+        return False
+
+    async def list_models(self) -> list[Model] | None:
+        return None
+
+    async def unregister_model(self, model_id: str) -> None:
+        pass
+
+    async def register_model(self, model: Model) -> Model:
+        llama_model = (
+            resolve_model(model.metadata["llama_model"])
+            if "llama_model" in model.metadata
+            else resolve_model(model.identifier)
+        )
+        if llama_model is None:
+            raise ValueError(
+                "Please make sure your llama_model in model metadata or model identifier is in Llama SKU list"
+            )
+
+        self.model_registry_helper = ModelRegistryHelper(
+            [
+                build_hf_repo_model_entry(
+                    llama_model.descriptor(),
+                    llama_model.core_model_id.value,
+                )
+            ],
+        )
+        model = await self.model_registry_helper.register_model(model)
+
+        if model.model_type == ModelType.embedding:
+            self._load_sentence_transformer_model(model.provider_resource_id)
+
+        # TODO: what is this?! you can't really specify skipping via model metadata
+        # kill this madness
+        if "skip_load" in model.metadata and model.metadata["skip_load"]:
+            return model
+
+        await self.load_model(model.identifier, llama_model)
+        return model
+
+    async def load_model(self, model_id, llama_model) -> None:
+        log.info(f"Loading model `{model_id}`")
+
+        builder_params = [self.config, model_id, llama_model]
+
+        if self.config.create_distributed_process_group:
+            self.generator = LlamaModelParallelGenerator(
+                model_parallel_size=self.config.model_parallel_size or llama_model.pth_file_count,
+                builder_fn=llama_builder_fn,
+                builder_params=builder_params,
+                formatter=(
+                    Llama4ChatFormat(Llama4Tokenizer.get_instance())
+                    if llama_model.model_family == ModelFamily.llama4
+                    else Llama3ChatFormat(Llama3Tokenizer.get_instance())
+                ),
+            )
+            self.generator.start()
+        else:
+            self.generator = llama_builder_fn(*builder_params)
+
+        self.model_id = model_id
+        self.llama_model = llama_model
+
+        log.info("Warming up...")
+        await self.openai_chat_completion(
+            model=model_id,
+            messages=[{"role": "user", "content": "Hi how are you?"}],
+            max_tokens=20,
+        )
+        log.info("Warmed up!")
+
+    def check_model(self, request) -> None:
+        if self.model_id is None or self.llama_model is None:
+            raise RuntimeError(
+                "No avaible model yet, please register your requested model or add your model in the resouces first"
+            )
+        elif request.model != self.model_id:
+            raise RuntimeError(f"Model mismatch: request model: {request.model} != loaded model: {self.model_id}")
+
+    async def openai_chat_completion(
+        self,
+        params: OpenAIChatCompletionRequestWithExtraBody,
+    ) -> OpenAIChatCompletion | AsyncIterator[OpenAIChatCompletionChunk]:
+        raise NotImplementedError("OpenAI chat completion not supported by meta-reference inference provider")

llama_stack/providers/inline/inference/meta_reference/model_parallel.py

@@ -0,0 +1,96 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from collections.abc import Callable, Generator
+from copy import deepcopy
+from functools import partial
+from typing import Any
+
+from llama_stack.models.llama.llama3.chat_format import ChatFormat as Llama3ChatFormat
+from llama_stack.models.llama.llama4.chat_format import ChatFormat as Llama4ChatFormat
+from llama_stack.providers.utils.inference.prompt_adapter import (
+    ChatCompletionRequestWithRawContent,
+    CompletionRequestWithRawContent,
+)
+
+from .parallel_utils import ModelParallelProcessGroup
+
+
+class ModelRunner:
+    def __init__(self, llama):
+        self.llama = llama
+
+    # the `task` object is the same that is sent to `ModelParallelProcessGroup.run_inference()`
+    def __call__(self, task: Any):
+        if task[0] == "chat_completion":
+            return self.llama.chat_completion(task[1])
+        else:
+            raise ValueError(f"Unexpected task type {task[0]}")
+
+
+def init_model_cb(
+    builder_fn: Callable,
+    params: list[Any],
+):
+    llama = builder_fn(*params)
+    return ModelRunner(llama)
+
+
+class LlamaModelParallelGenerator:
+    """
+    This abstraction exists so
+     - we can run model parallel code without needing to run the CLIs via torchrun
+     - this also enables use model parallel code within a notebook context.
+
+    A Context Manager is used to ensure that the model parallel process is started and stopped
+    correctly. This does make the ergonomics a little awkward, because it isn't immediately
+    clear at the callsite why we need to use a context manager.
+    """
+
+    def __init__(
+        self,
+        model_parallel_size: int,
+        builder_fn: Callable,
+        builder_params: list[Any],
+        formatter: Llama3ChatFormat | Llama4ChatFormat,
+    ):
+        self.model_parallel_size = model_parallel_size
+        self.builder_fn = builder_fn
+        self.builder_params = builder_params
+        self.formatter = formatter
+
+    def start(self):
+        self.__enter__()
+
+    def stop(self):
+        self.__exit__(None, None, None)
+
+    def __enter__(self):
+        self.group = ModelParallelProcessGroup(
+            self.model_parallel_size,
+            init_model_cb=partial(init_model_cb, self.builder_fn, self.builder_params),
+        )
+        self.group.start()
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        self.group.stop()
+
+    def completion(
+        self,
+        request_batch: list[CompletionRequestWithRawContent],
+    ) -> Generator:
+        req_obj = deepcopy(request_batch)
+        gen = self.group.run_inference(("completion", req_obj))
+        yield from gen
+
+    def chat_completion(
+        self,
+        request_batch: list[ChatCompletionRequestWithRawContent],
+    ) -> Generator:
+        req_obj = deepcopy(request_batch)
+        gen = self.group.run_inference(("chat_completion", req_obj))
+        yield from gen

llama_stack/providers/{impls/meta_reference/inference → inline/inference/meta_reference}/parallel_utils.py

@@ -4,42 +4,41 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+# Copyright (c) Meta Platforms, IAny, nc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import copy
 import json
 import multiprocessing
 import os
 import tempfile
 import time
 import uuid
+from collections.abc import Callable, Generator
 from enum import Enum
-from typing import Callable, Generator, List, Literal, Optional, Union
+from typing import Annotated, Literal
 
 import torch
-
 import zmq
-
 from fairscale.nn.model_parallel.initialize import (
     get_model_parallel_group,
     get_model_parallel_rank,
     get_model_parallel_src_rank,
 )
-
-from llama_models.llama3.api.datatypes import Message, ToolPromptFormat
-
 from pydantic import BaseModel, Field
+from torch.distributed.launcher.api import LaunchConfig, elastic_launch
 
-from torch.distributed.launcher.api import elastic_launch, LaunchConfig
-from typing_extensions import Annotated
-
-from .generation import TokenResult
-
+from llama_stack.log import get_logger
+from llama_stack.models.llama.datatypes import GenerationResult
+from llama_stack.providers.utils.inference.prompt_adapter import (
+    ChatCompletionRequestWithRawContent,
+    CompletionRequestWithRawContent,
+)
 
-class InferenceArgs(BaseModel):
-    messages: List[Message]
-    temperature: float
-    top_p: float
-    max_gen_len: int
-    logprobs: bool
-    tool_prompt_format: ToolPromptFormat
+log = get_logger(name=__name__, category="inference")
 
 
 class ProcessingMessageName(str, Enum):
@@ -53,59 +52,42 @@ class ProcessingMessageName(str, Enum):
 
 
 class ReadyRequest(BaseModel):
-    type: Literal[ProcessingMessageName.ready_request] = (
-        ProcessingMessageName.ready_request
-    )
+    type: Literal[ProcessingMessageName.ready_request] = ProcessingMessageName.ready_request
 
 
 class ReadyResponse(BaseModel):
-    type: Literal[ProcessingMessageName.ready_response] = (
-        ProcessingMessageName.ready_response
-    )
+    type: Literal[ProcessingMessageName.ready_response] = ProcessingMessageName.ready_response
 
 
 class EndSentinel(BaseModel):
-    type: Literal[ProcessingMessageName.end_sentinel] = (
-        ProcessingMessageName.end_sentinel
-    )
+    type: Literal[ProcessingMessageName.end_sentinel] = ProcessingMessageName.end_sentinel
 
 
 class CancelSentinel(BaseModel):
-    type: Literal[ProcessingMessageName.cancel_sentinel] = (
-        ProcessingMessageName.cancel_sentinel
-    )
+    type: Literal[ProcessingMessageName.cancel_sentinel] = ProcessingMessageName.cancel_sentinel
 
 
 class TaskRequest(BaseModel):
-    type: Literal[ProcessingMessageName.task_request] = (
-        ProcessingMessageName.task_request
-    )
-    task: InferenceArgs
+    type: Literal[ProcessingMessageName.task_request] = ProcessingMessageName.task_request
+    task: tuple[
+        str,
+        list[CompletionRequestWithRawContent] | list[ChatCompletionRequestWithRawContent],
+    ]
 
 
 class TaskResponse(BaseModel):
-    type: Literal[ProcessingMessageName.task_response] = (
-        ProcessingMessageName.task_response
-    )
-    result: TokenResult
+    type: Literal[ProcessingMessageName.task_response] = ProcessingMessageName.task_response
+    result: list[GenerationResult]
 
 
 class ExceptionResponse(BaseModel):
-    type: Literal[ProcessingMessageName.exception_response] = (
-        ProcessingMessageName.exception_response
-    )
+    type: Literal[ProcessingMessageName.exception_response] = ProcessingMessageName.exception_response
     error: str
 
 
-ProcessingMessage = Union[
-    ReadyRequest,
-    ReadyResponse,
-    EndSentinel,
-    CancelSentinel,
-    TaskRequest,
-    TaskResponse,
-    ExceptionResponse,
-]
+ProcessingMessage = (
+    ReadyRequest | ReadyResponse | EndSentinel | CancelSentinel | TaskRequest | TaskResponse | ExceptionResponse
+)
 
 
 class ProcessingMessageWrapper(BaseModel):
@@ -116,7 +98,7 @@ class ProcessingMessageWrapper(BaseModel):
 
 
 def mp_rank_0() -> bool:
-    return get_model_parallel_rank() == 0
+    return bool(get_model_parallel_rank() == 0)
 
 
 def encode_msg(msg: ProcessingMessage) -> bytes:
@@ -143,7 +125,7 @@ def retrieve_requests(reply_socket_url: str):
         reply_socket.send_multipart([client_id, encode_msg(obj)])
 
     while True:
-        tasks = [None]
+        tasks: list[ProcessingMessage | None] = [None]
         if mp_rank_0():
             client_id, maybe_task_json = maybe_get_work(reply_socket)
             if maybe_task_json is not None:
@@ -170,7 +152,7 @@ def retrieve_requests(reply_socket_url: str):
                     break
 
                 for obj in out:
-                    updates = [None]
+                    updates: list[ProcessingMessage | None] = [None]
                     if mp_rank_0():
                         _, update_json = maybe_get_work(reply_socket)
                         update = maybe_parse_message(update_json)
@@ -187,16 +169,14 @@ def retrieve_requests(reply_socket_url: str):
                         group=get_model_parallel_group(),
                     )
                     if isinstance(updates[0], CancelSentinel):
-                        print("quitting generation loop because request was cancelled")
+                        log.info("quitting generation loop because request was cancelled")
                         break
 
                 if mp_rank_0():
                     send_obj(EndSentinel())
             except Exception as e:
-                print(f"[debug] got exception {e}")
-                import traceback
+                log.exception("exception in generation loop")
 
-                traceback.print_exc()
                 if mp_rank_0():
                     send_obj(ExceptionResponse(error=str(e)))
 
@@ -217,20 +197,20 @@ def maybe_get_work(sock: zmq.Socket):
     return client_id, message
 
 
-def maybe_parse_message(maybe_json: Optional[str]) -> Optional[ProcessingMessage]:
+def maybe_parse_message(maybe_json: str | None) -> ProcessingMessage | None:
     if maybe_json is None:
         return None
     try:
         return parse_message(maybe_json)
     except json.JSONDecodeError:
         return None
-    except ValueError as e:
+    except ValueError:
         return None
 
 
 def parse_message(json_str: str) -> ProcessingMessage:
     data = json.loads(json_str)
-    return ProcessingMessageWrapper(**data).payload
+    return copy.deepcopy(ProcessingMessageWrapper(**data).payload)
 
 
 def worker_process_entrypoint(
@@ -248,15 +228,15 @@ def worker_process_entrypoint(
     while True:
         try:
             task = req_gen.send(result)
-            if isinstance(task, str) and task == _END_SENTINEL:
+            if isinstance(task, EndSentinel):
                 break
 
-            assert isinstance(task, TaskRequest)
+            assert isinstance(task, TaskRequest), task
             result = model(task.task)
         except StopIteration:
             break
 
-    print("[debug] worker process done")
+    log.info("[debug] worker process done")
 
 
 def launch_dist_group(
@@ -265,9 +245,6 @@ def launch_dist_group(
     init_model_cb: Callable,
     **kwargs,
 ) -> None:
-    id = uuid.uuid4().hex
-    dist_url = f"file:///tmp/llama3_{id}_{time.time()}"
-
     with tempfile.TemporaryDirectory() as tmpdir:
         # TODO: track workers and if they terminate, tell parent process about it so cleanup can happen
         launch_config = LaunchConfig(
@@ -301,7 +278,7 @@ def start_model_parallel_process(
 
     main_process_url = request_socket.getsockopt_string(zmq.LAST_ENDPOINT)
 
-    ctx = multiprocessing.get_context("fork")
+    ctx = multiprocessing.get_context("spawn")
     process = ctx.Process(
         target=launch_dist_group,
         args=(
@@ -316,8 +293,8 @@ def start_model_parallel_process(
     # wait until the model is loaded; rank 0 will send a message to indicate it's ready
 
     request_socket.send(encode_msg(ReadyRequest()))
-    response = request_socket.recv()
-    print("Loaded model...")
+    _response = request_socket.recv()
+    log.info("Loaded model...")
 
     return request_socket, process
 
@@ -349,12 +326,18 @@ class ModelParallelProcessGroup:
             self.process.join()
         self.started = False
 
-    def run_inference(self, inference_args: InferenceArgs) -> Generator:
+    def run_inference(
+        self,
+        req: tuple[
+            str,
+            list[CompletionRequestWithRawContent] | list[ChatCompletionRequestWithRawContent],
+        ],
+    ) -> Generator:
         assert not self.running, "inference already running"
 
         self.running = True
-        self.request_socket.send(encode_msg(TaskRequest(task=inference_args)))
         try:
+            self.request_socket.send(encode_msg(TaskRequest(task=req)))
             while True:
                 obj_json = self.request_socket.recv()
                 obj = parse_message(obj_json)
@@ -363,13 +346,13 @@ class ModelParallelProcessGroup:
                     break
 
                 if isinstance(obj, ExceptionResponse):
-                    print(f"[debug] got exception {obj.error}")
+                    log.error(f"[debug] got exception {obj.error}")
                     raise Exception(obj.error)
 
                 if isinstance(obj, TaskResponse):
                     yield obj.result
 
-        except GeneratorExit as e:
+        except GeneratorExit:
             self.request_socket.send(encode_msg(CancelSentinel()))
             while True:
                 obj_json = self.request_socket.send()

llama_stack/providers/inline/inference/sentence_transformers/__init__.py

@@ -0,0 +1,22 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Any
+
+from llama_stack.providers.inline.inference.sentence_transformers.config import (
+    SentenceTransformersInferenceConfig,
+)
+
+
+async def get_provider_impl(
+    config: SentenceTransformersInferenceConfig,
+    _deps: dict[str, Any],
+):
+    from .sentence_transformers import SentenceTransformersInferenceImpl
+
+    impl = SentenceTransformersInferenceImpl(config)
+    await impl.initialize()
+    return impl

llama_stack/providers/{impls/meta_reference/agents → inline/inference/sentence_transformers}/config.py

@@ -4,10 +4,12 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from pydantic import BaseModel
+from typing import Any
 
-from llama_stack.providers.utils.kvstore import KVStoreConfig
+from pydantic import BaseModel
 
 
-class MetaReferenceAgentsImplConfig(BaseModel):
-    persistence_store: KVStoreConfig
+class SentenceTransformersInferenceConfig(BaseModel):
+    @classmethod
+    def sample_run_config(cls, **kwargs) -> dict[str, Any]:
+        return {}

llama_stack/providers/inline/inference/sentence_transformers/sentence_transformers.py

@@ -0,0 +1,83 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from collections.abc import AsyncIterator
+
+from llama_stack.apis.inference import (
+    InferenceProvider,
+    OpenAIChatCompletionRequestWithExtraBody,
+    OpenAICompletionRequestWithExtraBody,
+)
+from llama_stack.apis.inference.inference import (
+    OpenAIChatCompletion,
+    OpenAIChatCompletionChunk,
+    OpenAICompletion,
+)
+from llama_stack.apis.models import ModelType
+from llama_stack.log import get_logger
+from llama_stack.providers.datatypes import Model, ModelsProtocolPrivate
+from llama_stack.providers.utils.inference.embedding_mixin import (
+    SentenceTransformerEmbeddingMixin,
+)
+from llama_stack.providers.utils.inference.openai_compat import (
+    OpenAIChatCompletionToLlamaStackMixin,
+)
+
+from .config import SentenceTransformersInferenceConfig
+
+log = get_logger(name=__name__, category="inference")
+
+
+class SentenceTransformersInferenceImpl(
+    OpenAIChatCompletionToLlamaStackMixin,
+    SentenceTransformerEmbeddingMixin,
+    InferenceProvider,
+    ModelsProtocolPrivate,
+):
+    __provider_id__: str
+
+    def __init__(self, config: SentenceTransformersInferenceConfig) -> None:
+        self.config = config
+
+    async def initialize(self) -> None:
+        pass
+
+    async def shutdown(self) -> None:
+        pass
+
+    async def should_refresh_models(self) -> bool:
+        return False
+
+    async def list_models(self) -> list[Model] | None:
+        return [
+            Model(
+                identifier="nomic-ai/nomic-embed-text-v1.5",
+                provider_resource_id="nomic-ai/nomic-embed-text-v1.5",
+                provider_id=self.__provider_id__,
+                metadata={
+                    "embedding_dimension": 768,
+                },
+                model_type=ModelType.embedding,
+            ),
+        ]
+
+    async def register_model(self, model: Model) -> Model:
+        return model
+
+    async def unregister_model(self, model_id: str) -> None:
+        pass
+
+    async def openai_completion(
+        self,
+        params: OpenAICompletionRequestWithExtraBody,
+    ) -> OpenAICompletion:
+        raise NotImplementedError("OpenAI completion not supported by sentence transformers provider")
+
+    async def openai_chat_completion(
+        self,
+        params: OpenAIChatCompletionRequestWithExtraBody,
+    ) -> OpenAIChatCompletion | AsyncIterator[OpenAIChatCompletionChunk]:
+        raise NotImplementedError("OpenAI chat completion not supported by sentence transformers provider")

llama_stack/providers/inline/post_training/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.

llama_stack/providers/inline/post_training/common/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.