llama-stack 0.4.4__py3-none-any.whl → 0.5.0__py3-none-any.whl

llama_stack/providers/inline/agents/meta_reference/responses/types.py

@@ -12,6 +12,7 @@ from pydantic import BaseModel
 
 from llama_stack_api import (
     OpenAIChatCompletionToolCall,
+    OpenAIFinishReason,
     OpenAIMessageParam,
     OpenAIResponseFormatParam,
     OpenAIResponseInput,
@@ -52,7 +53,7 @@ class ChatCompletionResult:
     tool_calls: dict[int, OpenAIChatCompletionToolCall]
     created: int
     model: str
-    finish_reason: str
+    finish_reason: OpenAIFinishReason
     message_item_id: str  # For streaming events
     tool_call_item_ids: dict[int, str]  # For streaming events
     content_part_emitted: bool  # Tracking state

llama_stack/providers/inline/agents/meta_reference/responses/utils.py

@@ -53,6 +53,7 @@ from llama_stack_api import (
     OpenAIToolMessageParam,
     OpenAIUserMessageParam,
     ResponseGuardrailSpec,
+    RunModerationRequest,
     Safety,
 )
 
@@ -468,7 +469,9 @@ async def run_guardrails(safety_api: Safety | None, messages: str, guardrail_ids
         else:
             raise ValueError(f"No shield found with identifier '{guardrail_id}'")
 
-    guardrail_tasks = [safety_api.run_moderation(messages, model=model_id) for model_id in model_ids]
+    guardrail_tasks = [
+        safety_api.run_moderation(RunModerationRequest(input=messages, model=model_id)) for model_id in model_ids
+    ]
     responses = await asyncio.gather(*guardrail_tasks)
 
     for response in responses:

llama_stack/providers/inline/agents/meta_reference/safety.py

@@ -7,7 +7,7 @@
 import asyncio
 
 from llama_stack.log import get_logger
-from llama_stack_api import OpenAIMessageParam, Safety, SafetyViolation, ViolationLevel
+from llama_stack_api import OpenAIMessageParam, RunShieldRequest, Safety, SafetyViolation, ViolationLevel
 
 log = get_logger(name=__name__, category="agents::meta_reference")
 
@@ -32,7 +32,7 @@ class ShieldRunnerMixin:
     async def run_multiple_shields(self, messages: list[OpenAIMessageParam], identifiers: list[str]) -> None:
         responses = await asyncio.gather(
             *[
-                self.safety_api.run_shield(shield_id=identifier, messages=messages, params={})
+                self.safety_api.run_shield(RunShieldRequest(shield_id=identifier, messages=messages))
                 for identifier in identifiers
             ]
         )

llama_stack/providers/inline/batches/reference/batches.py

@@ -23,6 +23,7 @@ from llama_stack_api import (
     BatchObject,
     ConflictError,
     Files,
+    GetModelRequest,
     Inference,
     ListBatchesResponse,
     Models,
@@ -485,7 +486,7 @@ class ReferenceBatchesImpl(Batches):
 
                         if "model" in request_body and isinstance(request_body["model"], str):
                             try:
-                                await self.models_api.get_model(request_body["model"])
+                                await self.models_api.get_model(GetModelRequest(model_id=request_body["model"]))
                             except Exception:
                                 errors.append(
                                     BatchError(

llama_stack/providers/inline/eval/meta_reference/eval.py

@@ -13,19 +13,25 @@ from llama_stack.providers.utils.common.data_schema_validator import ColumnName
 from llama_stack_api import (
     Agents,
     Benchmark,
-    BenchmarkConfig,
     BenchmarksProtocolPrivate,
     DatasetIO,
     Datasets,
     Eval,
     EvaluateResponse,
+    EvaluateRowsRequest,
     Inference,
+    IterRowsRequest,
     Job,
+    JobCancelRequest,
+    JobResultRequest,
     JobStatus,
+    JobStatusRequest,
     OpenAIChatCompletionRequestWithExtraBody,
     OpenAICompletionRequestWithExtraBody,
     OpenAISystemMessageParam,
     OpenAIUserMessageParam,
+    RunEvalRequest,
+    ScoreRequest,
     Scoring,
 )
 
@@ -90,10 +96,9 @@ class MetaReferenceEvalImpl(
 
     async def run_eval(
         self,
-        benchmark_id: str,
-        benchmark_config: BenchmarkConfig,
+        request: RunEvalRequest,
     ) -> Job:
-        task_def = self.benchmarks[benchmark_id]
+        task_def = self.benchmarks[request.benchmark_id]
         dataset_id = task_def.dataset_id
         scoring_functions = task_def.scoring_functions
 
@@ -101,15 +106,18 @@ class MetaReferenceEvalImpl(
         # dataset_def = await self.datasets_api.get_dataset(dataset_id=dataset_id)
 
         all_rows = await self.datasetio_api.iterrows(
-            dataset_id=dataset_id,
-            limit=(-1 if benchmark_config.num_examples is None else benchmark_config.num_examples),
+            IterRowsRequest(
+                dataset_id=dataset_id,
+                limit=(-1 if request.benchmark_config.num_examples is None else request.benchmark_config.num_examples),
+            )
         )
-        res = await self.evaluate_rows(
-            benchmark_id=benchmark_id,
+        eval_rows_request = EvaluateRowsRequest(
+            benchmark_id=request.benchmark_id,
             input_rows=all_rows.data,
             scoring_functions=scoring_functions,
-            benchmark_config=benchmark_config,
+            benchmark_config=request.benchmark_config,
         )
+        res = await self.evaluate_rows(eval_rows_request)
 
         # TODO: currently needs to wait for generation before returning
         # need job scheduler queue (ray/celery) w/ jobs api
@@ -118,9 +126,9 @@ class MetaReferenceEvalImpl(
         return Job(job_id=job_id, status=JobStatus.completed)
 
     async def _run_model_generation(
-        self, input_rows: list[dict[str, Any]], benchmark_config: BenchmarkConfig
+        self, input_rows: list[dict[str, Any]], request: EvaluateRowsRequest
     ) -> list[dict[str, Any]]:
-        candidate = benchmark_config.eval_candidate
+        candidate = request.benchmark_config.eval_candidate
         assert candidate.sampling_params.max_tokens is not None, "SamplingParams.max_tokens must be provided"
         sampling_params = {"max_tokens": candidate.sampling_params.max_tokens}
 
@@ -165,50 +173,50 @@ class MetaReferenceEvalImpl(
 
     async def evaluate_rows(
         self,
-        benchmark_id: str,
-        input_rows: list[dict[str, Any]],
-        scoring_functions: list[str],
-        benchmark_config: BenchmarkConfig,
+        request: EvaluateRowsRequest,
     ) -> EvaluateResponse:
-        candidate = benchmark_config.eval_candidate
+        candidate = request.benchmark_config.eval_candidate
         # Agent evaluation removed
         if candidate.type == "model":
-            generations = await self._run_model_generation(input_rows, benchmark_config)
+            generations = await self._run_model_generation(request.input_rows, request)
         else:
             raise ValueError(f"Invalid candidate type: {candidate.type}")
 
         # scoring with generated_answer
         score_input_rows = [
-            input_r | generated_r for input_r, generated_r in zip(input_rows, generations, strict=False)
+            input_r | generated_r for input_r, generated_r in zip(request.input_rows, generations, strict=False)
         ]
 
-        if benchmark_config.scoring_params is not None:
+        if request.benchmark_config.scoring_params is not None:
             scoring_functions_dict = {
-                scoring_fn_id: benchmark_config.scoring_params.get(scoring_fn_id, None)
-                for scoring_fn_id in scoring_functions
+                scoring_fn_id: request.benchmark_config.scoring_params.get(scoring_fn_id, None)
+                for scoring_fn_id in request.scoring_functions
             }
         else:
-            scoring_functions_dict = dict.fromkeys(scoring_functions)
+            scoring_functions_dict = dict.fromkeys(request.scoring_functions)
 
-        score_response = await self.scoring_api.score(
-            input_rows=score_input_rows, scoring_functions=scoring_functions_dict
+        score_request = ScoreRequest(
+            input_rows=score_input_rows,
+            scoring_functions=scoring_functions_dict,
         )
+        score_response = await self.scoring_api.score(score_request)
 
         return EvaluateResponse(generations=generations, scores=score_response.results)
 
-    async def job_status(self, benchmark_id: str, job_id: str) -> Job:
-        if job_id in self.jobs:
-            return Job(job_id=job_id, status=JobStatus.completed)
+    async def job_status(self, request: JobStatusRequest) -> Job:
+        if request.job_id in self.jobs:
+            return Job(job_id=request.job_id, status=JobStatus.completed)
 
-        raise ValueError(f"Job {job_id} not found")
+        raise ValueError(f"Job {request.job_id} not found")
 
-    async def job_cancel(self, benchmark_id: str, job_id: str) -> None:
+    async def job_cancel(self, request: JobCancelRequest) -> None:
         raise NotImplementedError("Job cancel is not implemented yet")
 
-    async def job_result(self, benchmark_id: str, job_id: str) -> EvaluateResponse:
-        job = await self.job_status(benchmark_id, job_id)
+    async def job_result(self, request: JobResultRequest) -> EvaluateResponse:
+        job_status_request = JobStatusRequest(benchmark_id=request.benchmark_id, job_id=request.job_id)
+        job = await self.job_status(job_status_request)
         status = job.status
         if not status or status != JobStatus.completed:
             raise ValueError(f"Job is not completed, Status: {status.value}")
 
-        return self.jobs[job_id]
+        return self.jobs[request.job_id]

llama_stack/providers/inline/post_training/huggingface/post_training.py

@@ -12,17 +12,19 @@ from llama_stack.providers.inline.post_training.huggingface.config import (
 from llama_stack.providers.utils.scheduler import JobArtifact, Scheduler
 from llama_stack.providers.utils.scheduler import JobStatus as SchedulerJobStatus
 from llama_stack_api import (
-    AlgorithmConfig,
+    CancelTrainingJobRequest,
     Checkpoint,
     DatasetIO,
     Datasets,
-    DPOAlignmentConfig,
+    GetTrainingJobArtifactsRequest,
+    GetTrainingJobStatusRequest,
     JobStatus,
     ListPostTrainingJobsResponse,
     PostTrainingJob,
     PostTrainingJobArtifactsResponse,
     PostTrainingJobStatusResponse,
-    TrainingConfig,
+    PreferenceOptimizeRequest,
+    SupervisedFineTuneRequest,
 )
 
 
@@ -69,13 +71,7 @@ class HuggingFacePostTrainingImpl:
 
     async def supervised_fine_tune(
         self,
-        job_uuid: str,
-        training_config: TrainingConfig,
-        hyperparam_search_config: dict[str, Any],
-        logger_config: dict[str, Any],
-        model: str,
-        checkpoint_dir: str | None = None,
-        algorithm_config: AlgorithmConfig | None = None,
+        request: SupervisedFineTuneRequest,
     ) -> PostTrainingJob:
         async def handler(on_log_message_cb, on_status_change_cb, on_artifact_collected_cb):
             from llama_stack.providers.inline.post_training.huggingface.recipes.finetune_single_device import (
@@ -85,17 +81,17 @@ class HuggingFacePostTrainingImpl:
             on_log_message_cb("Starting HF finetuning")
 
             recipe = HFFinetuningSingleDevice(
-                job_uuid=job_uuid,
+                job_uuid=request.job_uuid,
                 datasetio_api=self.datasetio_api,
                 datasets_api=self.datasets_api,
             )
 
             resources_allocated, checkpoints = await recipe.train(
-                model=model,
-                output_dir=checkpoint_dir,
-                job_uuid=job_uuid,
-                lora_config=algorithm_config,
-                config=training_config,
+                model=request.model,
+                output_dir=request.checkpoint_dir,
+                job_uuid=request.job_uuid,
+                lora_config=request.algorithm_config,
+                config=request.training_config,
                 provider_config=self.config,
             )
 
@@ -108,17 +104,12 @@ class HuggingFacePostTrainingImpl:
             on_status_change_cb(SchedulerJobStatus.completed)
             on_log_message_cb("HF finetuning completed")
 
-        job_uuid = self._scheduler.schedule(_JOB_TYPE_SUPERVISED_FINE_TUNE, job_uuid, handler)
+        job_uuid = self._scheduler.schedule(_JOB_TYPE_SUPERVISED_FINE_TUNE, request.job_uuid, handler)
         return PostTrainingJob(job_uuid=job_uuid)
 
     async def preference_optimize(
         self,
-        job_uuid: str,
-        finetuned_model: str,
-        algorithm_config: DPOAlignmentConfig,
-        training_config: TrainingConfig,
-        hyperparam_search_config: dict[str, Any],
-        logger_config: dict[str, Any],
+        request: PreferenceOptimizeRequest,
     ) -> PostTrainingJob:
         async def handler(on_log_message_cb, on_status_change_cb, on_artifact_collected_cb):
             from llama_stack.providers.inline.post_training.huggingface.recipes.finetune_single_device_dpo import (
@@ -128,17 +119,17 @@ class HuggingFacePostTrainingImpl:
             on_log_message_cb("Starting HF DPO alignment")
 
             recipe = HFDPOAlignmentSingleDevice(
-                job_uuid=job_uuid,
+                job_uuid=request.job_uuid,
                 datasetio_api=self.datasetio_api,
                 datasets_api=self.datasets_api,
             )
 
             resources_allocated, checkpoints = await recipe.train(
-                model=finetuned_model,
-                output_dir=f"{self.config.dpo_output_dir}/{job_uuid}",
-                job_uuid=job_uuid,
-                dpo_config=algorithm_config,
-                config=training_config,
+                model=request.finetuned_model,
+                output_dir=f"{self.config.dpo_output_dir}/{request.job_uuid}",
+                job_uuid=request.job_uuid,
+                dpo_config=request.algorithm_config,
+                config=request.training_config,
                 provider_config=self.config,
             )
 
@@ -153,7 +144,7 @@ class HuggingFacePostTrainingImpl:
             on_status_change_cb(SchedulerJobStatus.completed)
             on_log_message_cb("HF DPO alignment completed")
 
-        job_uuid = self._scheduler.schedule(_JOB_TYPE_DPO_TRAINING, job_uuid, handler)
+        job_uuid = self._scheduler.schedule(_JOB_TYPE_DPO_TRAINING, request.job_uuid, handler)
         return PostTrainingJob(job_uuid=job_uuid)
 
     @staticmethod
@@ -169,8 +160,10 @@ class HuggingFacePostTrainingImpl:
         data = cls._get_artifacts_metadata_by_type(job, TrainingArtifactType.RESOURCES_STATS.value)
         return data[0] if data else None
 
-    async def get_training_job_status(self, job_uuid: str) -> PostTrainingJobStatusResponse | None:
-        job = self._scheduler.get_job(job_uuid)
+    async def get_training_job_status(
+        self, request: GetTrainingJobStatusRequest
+    ) -> PostTrainingJobStatusResponse | None:
+        job = self._scheduler.get_job(request.job_uuid)
 
         match job.status:
             # TODO: Add support for other statuses to API
@@ -186,7 +179,7 @@ class HuggingFacePostTrainingImpl:
                 raise NotImplementedError()
 
         return PostTrainingJobStatusResponse(
-            job_uuid=job_uuid,
+            job_uuid=request.job_uuid,
             status=status,
             scheduled_at=job.scheduled_at,
             started_at=job.started_at,
@@ -195,12 +188,14 @@ class HuggingFacePostTrainingImpl:
             resources_allocated=self._get_resources_allocated(job),
         )
 
-    async def cancel_training_job(self, job_uuid: str) -> None:
-        self._scheduler.cancel(job_uuid)
+    async def cancel_training_job(self, request: CancelTrainingJobRequest) -> None:
+        self._scheduler.cancel(request.job_uuid)
 
-    async def get_training_job_artifacts(self, job_uuid: str) -> PostTrainingJobArtifactsResponse | None:
-        job = self._scheduler.get_job(job_uuid)
-        return PostTrainingJobArtifactsResponse(job_uuid=job_uuid, checkpoints=self._get_checkpoints(job))
+    async def get_training_job_artifacts(
+        self, request: GetTrainingJobArtifactsRequest
+    ) -> PostTrainingJobArtifactsResponse | None:
+        job = self._scheduler.get_job(request.job_uuid)
+        return PostTrainingJobArtifactsResponse(job_uuid=request.job_uuid, checkpoints=self._get_checkpoints(job))
 
     async def get_training_jobs(self) -> ListPostTrainingJobsResponse:
         return ListPostTrainingJobsResponse(

llama_stack/providers/inline/post_training/huggingface/utils.py

@@ -16,7 +16,7 @@ import torch
 from datasets import Dataset
 from transformers import AutoConfig, AutoModelForCausalLM
 
-from llama_stack_api import Checkpoint, DatasetIO, TrainingConfig
+from llama_stack_api import Checkpoint, DatasetIO, IterRowsRequest, TrainingConfig
 
 if TYPE_CHECKING:
     from transformers import PretrainedConfig
@@ -135,10 +135,7 @@ def setup_torch_device(device_str: str) -> torch.device:
 async def load_rows_from_dataset(datasetio_api: DatasetIO, dataset_id: str) -> list[dict[str, Any]]:
     """Load dataset from llama stack dataset provider"""
     try:
-        all_rows = await datasetio_api.iterrows(
-            dataset_id=dataset_id,
-            limit=-1,
-        )
+        all_rows = await datasetio_api.iterrows(IterRowsRequest(dataset_id=dataset_id, limit=-1))
         if not isinstance(all_rows.data, list):
             raise RuntimeError("Expected dataset data to be a list")
         return all_rows.data

llama_stack/providers/inline/post_training/torchtune/common/utils.py

@@ -22,7 +22,6 @@ from torchtune.models.llama3_2 import lora_llama3_2_3b
 from torchtune.modules.transforms import Transform
 
 from llama_stack.models.llama.sku_list import resolve_model
-from llama_stack.models.llama.sku_types import Model
 from llama_stack_api import DatasetFormat
 
 BuildLoraModelCallable = Callable[..., torch.nn.Module]
@@ -54,18 +53,17 @@ DATA_FORMATS: dict[str, Transform] = {
 }
 
 
-def _validate_model_id(model_id: str) -> Model:
+def _validate_model_id(model_id: str) -> str:
     model = resolve_model(model_id)
     if model is None or model.core_model_id.value not in MODEL_CONFIGS:
         raise ValueError(f"Model {model_id} is not supported.")
-    return model
+    return model.core_model_id.value
 
 
 async def get_model_definition(
     model_id: str,
 ) -> BuildLoraModelCallable:
-    model = _validate_model_id(model_id)
-    model_config = MODEL_CONFIGS[model.core_model_id.value]
+    model_config = MODEL_CONFIGS[_validate_model_id(model_id)]
     if not hasattr(model_config, "model_definition"):
         raise ValueError(f"Model {model_id} does not have model definition.")
     return model_config.model_definition
@@ -74,8 +72,7 @@ async def get_model_definition(
 async def get_tokenizer_type(
     model_id: str,
 ) -> BuildTokenizerCallable:
-    model = _validate_model_id(model_id)
-    model_config = MODEL_CONFIGS[model.core_model_id.value]
+    model_config = MODEL_CONFIGS[_validate_model_id(model_id)]
     if not hasattr(model_config, "tokenizer_type"):
         raise ValueError(f"Model {model_id} does not have tokenizer_type.")
     return model_config.tokenizer_type
@@ -88,8 +85,7 @@ async def get_checkpointer_model_type(
     checkpointer model type is used in checkpointer for some special treatment on some specific model types
     For example, llama3.2 model tied weights (https://github.com/pytorch/torchtune/blob/main/torchtune/training/checkpointing/_checkpointer.py#L1041)
     """
-    model = _validate_model_id(model_id)
-    model_config = MODEL_CONFIGS[model.core_model_id.value]
+    model_config = MODEL_CONFIGS[_validate_model_id(model_id)]
     if not hasattr(model_config, "checkpoint_type"):
         raise ValueError(f"Model {model_id} does not have checkpoint_type.")
     return model_config.checkpoint_type

llama_stack/providers/inline/post_training/torchtune/post_training.py

@@ -12,18 +12,20 @@ from llama_stack.providers.inline.post_training.torchtune.config import (
 from llama_stack.providers.utils.scheduler import JobArtifact, Scheduler
 from llama_stack.providers.utils.scheduler import JobStatus as SchedulerJobStatus
 from llama_stack_api import (
-    AlgorithmConfig,
+    CancelTrainingJobRequest,
     Checkpoint,
     DatasetIO,
     Datasets,
-    DPOAlignmentConfig,
+    GetTrainingJobArtifactsRequest,
+    GetTrainingJobStatusRequest,
     JobStatus,
     ListPostTrainingJobsResponse,
     LoraFinetuningConfig,
     PostTrainingJob,
     PostTrainingJobArtifactsResponse,
     PostTrainingJobStatusResponse,
-    TrainingConfig,
+    PreferenceOptimizeRequest,
+    SupervisedFineTuneRequest,
 )
 
 
@@ -69,15 +71,9 @@ class TorchtunePostTrainingImpl:
 
     async def supervised_fine_tune(
         self,
-        job_uuid: str,
-        training_config: TrainingConfig,
-        hyperparam_search_config: dict[str, Any],
-        logger_config: dict[str, Any],
-        model: str,
-        checkpoint_dir: str | None,
-        algorithm_config: AlgorithmConfig | None,
+        request: SupervisedFineTuneRequest,
     ) -> PostTrainingJob:
-        if isinstance(algorithm_config, LoraFinetuningConfig):
+        if isinstance(request.algorithm_config, LoraFinetuningConfig):
 
             async def handler(on_log_message_cb, on_status_change_cb, on_artifact_collected_cb):
                 from llama_stack.providers.inline.post_training.torchtune.recipes.lora_finetuning_single_device import (
@@ -88,13 +84,13 @@ class TorchtunePostTrainingImpl:
 
                 recipe = LoraFinetuningSingleDevice(
                     self.config,
-                    job_uuid,
-                    training_config,
-                    hyperparam_search_config,
-                    logger_config,
-                    model,
-                    checkpoint_dir,
-                    algorithm_config,
+                    request.job_uuid,
+                    request.training_config,
+                    request.hyperparam_search_config,
+                    request.logger_config,
+                    request.model,
+                    request.checkpoint_dir,
+                    request.algorithm_config,
                     self.datasetio_api,
                     self.datasets_api,
                 )
@@ -112,17 +108,12 @@ class TorchtunePostTrainingImpl:
         else:
             raise NotImplementedError()
 
-        job_uuid = self._scheduler.schedule(_JOB_TYPE_SUPERVISED_FINE_TUNE, job_uuid, handler)
+        job_uuid = self._scheduler.schedule(_JOB_TYPE_SUPERVISED_FINE_TUNE, request.job_uuid, handler)
         return PostTrainingJob(job_uuid=job_uuid)
 
     async def preference_optimize(
         self,
-        job_uuid: str,
-        finetuned_model: str,
-        algorithm_config: DPOAlignmentConfig,
-        training_config: TrainingConfig,
-        hyperparam_search_config: dict[str, Any],
-        logger_config: dict[str, Any],
+        request: PreferenceOptimizeRequest,
     ) -> PostTrainingJob:
         raise NotImplementedError()
 
@@ -144,8 +135,10 @@ class TorchtunePostTrainingImpl:
         data = cls._get_artifacts_metadata_by_type(job, TrainingArtifactType.RESOURCES_STATS.value)
         return data[0] if data else None
 
-    async def get_training_job_status(self, job_uuid: str) -> PostTrainingJobStatusResponse | None:
-        job = self._scheduler.get_job(job_uuid)
+    async def get_training_job_status(
+        self, request: GetTrainingJobStatusRequest
+    ) -> PostTrainingJobStatusResponse | None:
+        job = self._scheduler.get_job(request.job_uuid)
 
         match job.status:
             # TODO: Add support for other statuses to API
@@ -161,7 +154,7 @@ class TorchtunePostTrainingImpl:
                 raise NotImplementedError()
 
         return PostTrainingJobStatusResponse(
-            job_uuid=job_uuid,
+            job_uuid=request.job_uuid,
             status=status,
             scheduled_at=job.scheduled_at,
             started_at=job.started_at,
@@ -170,9 +163,11 @@ class TorchtunePostTrainingImpl:
             resources_allocated=self._get_resources_allocated(job),
         )
 
-    async def cancel_training_job(self, job_uuid: str) -> None:
-        self._scheduler.cancel(job_uuid)
+    async def cancel_training_job(self, request: CancelTrainingJobRequest) -> None:
+        self._scheduler.cancel(request.job_uuid)
 
-    async def get_training_job_artifacts(self, job_uuid: str) -> PostTrainingJobArtifactsResponse | None:
-        job = self._scheduler.get_job(job_uuid)
-        return PostTrainingJobArtifactsResponse(job_uuid=job_uuid, checkpoints=self._get_checkpoints(job))
+    async def get_training_job_artifacts(
+        self, request: GetTrainingJobArtifactsRequest
+    ) -> PostTrainingJobArtifactsResponse | None:
+        job = self._scheduler.get_job(request.job_uuid)
+        return PostTrainingJobArtifactsResponse(job_uuid=request.job_uuid, checkpoints=self._get_checkpoints(job))

llama_stack/providers/inline/post_training/torchtune/recipes/lora_finetuning_single_device.py

@@ -50,6 +50,7 @@ from llama_stack_api import (
     DataConfig,
     DatasetIO,
     Datasets,
+    IterRowsRequest,
     LoraFinetuningConfig,
     OptimizerConfig,
     PostTrainingMetric,
@@ -334,10 +335,7 @@ class LoraFinetuningSingleDevice:
         batch_size: int,
     ) -> tuple[DistributedSampler, DataLoader]:
         async def fetch_rows(dataset_id: str):
-            return await self.datasetio_api.iterrows(
-                dataset_id=dataset_id,
-                limit=-1,
-            )
+            return await self.datasetio_api.iterrows(IterRowsRequest(dataset_id=dataset_id, limit=-1))
 
         all_rows = await fetch_rows(dataset_id)
         rows = all_rows.data

llama_stack/providers/inline/safety/code_scanner/code_scanner.py

@@ -5,7 +5,7 @@
 # the root directory of this source tree.
 
 import uuid
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from codeshield.cs import CodeShieldScanResult
@@ -15,9 +15,11 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
     interleaved_content_as_str,
 )
 from llama_stack_api import (
+    GetShieldRequest,
     ModerationObject,
     ModerationObjectResults,
-    OpenAIMessageParam,
+    RunModerationRequest,
+    RunShieldRequest,
     RunShieldResponse,
     Safety,
     SafetyViolation,
@@ -51,19 +53,14 @@ class MetaReferenceCodeScannerSafetyImpl(Safety):
                 f"Unsupported Code Scanner ID: {shield.provider_resource_id}. Allowed IDs: {ALLOWED_CODE_SCANNER_MODEL_IDS}"
             )
 
-    async def run_shield(
-        self,
-        shield_id: str,
-        messages: list[OpenAIMessageParam],
-        params: dict[str, Any] = None,
-    ) -> RunShieldResponse:
-        shield = await self.shield_store.get_shield(shield_id)
+    async def run_shield(self, request: RunShieldRequest) -> RunShieldResponse:
+        shield = await self.shield_store.get_shield(GetShieldRequest(identifier=request.shield_id))
         if not shield:
-            raise ValueError(f"Shield {shield_id} not found")
+            raise ValueError(f"Shield {request.shield_id} not found")
 
         from codeshield.cs import CodeShield
 
-        text = "\n".join([interleaved_content_as_str(m.content) for m in messages])
+        text = "\n".join([interleaved_content_as_str(m.content) for m in request.messages])
         log.info(f"Running CodeScannerShield on {text[50:]}")
         result = await CodeShield.scan_code(text)
 
@@ -102,11 +99,11 @@ class MetaReferenceCodeScannerSafetyImpl(Safety):
             metadata=metadata,
         )
 
-    async def run_moderation(self, input: str | list[str], model: str | None = None) -> ModerationObject:
-        if model is None:
+    async def run_moderation(self, request: RunModerationRequest) -> ModerationObject:
+        if request.model is None:
             raise ValueError("Code scanner moderation requires a model identifier.")
 
-        inputs = input if isinstance(input, list) else [input]
+        inputs = request.input if isinstance(request.input, list) else [request.input]
         results = []
 
         from codeshield.cs import CodeShield
@@ -129,4 +126,4 @@ class MetaReferenceCodeScannerSafetyImpl(Safety):
                 )
             results.append(moderation_result)
 
-        return ModerationObject(id=str(uuid.uuid4()), model=model, results=results)
+        return ModerationObject(id=str(uuid.uuid4()), model=request.model, results=results)