phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-06-28 02:53:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Update Telemetry API so OpenAPI generation can work (#640)

Browse source 
We cannot use recursive types because not only does our OpenAPI
generator not like them, even if it did, it is not easy for all client
languages to automatically construct proper APIs (especially considering
garbage collection) around them. For now, we can return a `Dict[str,
SpanWithStatus]` instead of `SpanWithChildren` and rely on the client to
reconstruct the tree.

Also fixed a super subtle issue with the OpenAPI generation process
(monkey-patching of json_schema_type wasn't working because of import
reordering.)
This commit is contained in:
Ashwin Bharambe 2024-12-16 13:00:14 -08:00 • committed by GitHub
parent 78e2bfbe7a
commit 2e5bfcd42a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 349 additions and 473 deletions
Download patch file Download diff file Expand all files Collapse all files

1
.gitignore vendored
View file

@ -18,3 +18,4 @@ Package.resolved
.vscode .vscode
_build _build
docs/src docs/src
pyrightconfig.json

10
docs/openapi_generator/generate.py
View file

@ -18,10 +18,6 @@ import yaml
from llama_models import schema_utils from llama_models import schema_utils
from .pyopenapi.options import Options
from .pyopenapi.specification import Info, Server
from .pyopenapi.utility import Specification
# We do some monkey-patching to ensure our definitions only use the minimal # We do some monkey-patching to ensure our definitions only use the minimal
# (json_schema_type, webmethod) definitions from the llama_models package. For # (json_schema_type, webmethod) definitions from the llama_models package. For
# generation though, we need the full definitions and implementations from the # generation though, we need the full definitions and implementations from the
@ -31,11 +27,13 @@ from .strong_typing.schema import json_schema_type
schema_utils.json_schema_type = json_schema_type schema_utils.json_schema_type = json_schema_type
# this line needs to be here to ensure json_schema_type has been altered before
# the imports use the annotation
from llama_stack.apis.version import LLAMA_STACK_API_VERSION # noqa: E402 from llama_stack.apis.version import LLAMA_STACK_API_VERSION # noqa: E402
from llama_stack.distribution.stack import LlamaStack # noqa: E402 from llama_stack.distribution.stack import LlamaStack # noqa: E402
from .pyopenapi.options import Options # noqa: E402
from .pyopenapi.specification import Info, Server # noqa: E402
from .pyopenapi.utility import Specification # noqa: E402
def main(output_dir: str): def main(output_dir: str):
output_dir = Path(output_dir) output_dir = Path(output_dir)

430
docs/resources/llama-stack-spec.html
View file

@ -1067,7 +1067,10 @@
"content": { "content": {
"application/json": { "application/json": {
"schema": { "schema": {
"$ref": "#/components/schemas/SpanWithChildren" "type": "object",
"additionalProperties": {
"$ref": "#/components/schemas/SpanWithStatus"
}
} }
} }
} }
@ -1123,46 +1126,15 @@
"content": { "content": {
"application/json": { "application/json": {
"schema": { "schema": {
"oneOf": [
{
"$ref": "#/components/schemas/PostTrainingJobArtifactsResponse" "$ref": "#/components/schemas/PostTrainingJobArtifactsResponse"
}
}
}
}
},
"tags": [
"PostTraining (Coming Soon)"
],
"parameters": [
{
"name": "job_uuid",
"in": "query",
"required": true,
"schema": {
"type": "string"
}
}, },
{ {
"name": "X-LlamaStack-ProviderData", "type": "null"
"in": "header",
"description": "JSON-encoded provider data which will be made available to the adapter servicing the API",
"required": false,
"schema": {
"type": "string"
}
} }
] ]
} }
},
"/alpha/post-training/job/logs": {
"get": {
"responses": {
"200": {
"description": "OK",
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/PostTrainingJobLogStream"
}
} }
} }
} }
@ -1199,7 +1171,14 @@
"content": { "content": {
"application/json": { "application/json": {
"schema": { "schema": {
"oneOf": [
{
"$ref": "#/components/schemas/PostTrainingJobStatusResponse" "$ref": "#/components/schemas/PostTrainingJobStatusResponse"
},
{
"type": "null"
}
]
} }
} }
} }
@ -5459,6 +5438,10 @@
"chunk_size_in_tokens": { "chunk_size_in_tokens": {
"type": "integer" "type": "integer"
}, },
"embedding_dimension": {
"type": "integer",
"default": 384
},
"overlap_size_in_tokens": { "overlap_size_in_tokens": {
"type": "integer" "type": "integer"
} }
@ -5807,6 +5790,10 @@
} }
] ]
} }
},
"model_type": {
"$ref": "#/components/schemas/ModelType",
"default": "llm"
} }
}, },
"additionalProperties": false, "additionalProperties": false,
@ -5815,7 +5802,15 @@
"provider_resource_id", "provider_resource_id",
"provider_id", "provider_id",
"type", "type",
"metadata" "metadata",
"model_type"
]
},
"ModelType": {
"type": "string",
"enum": [
"llm",
"embedding"
] ]
}, },
"PaginatedRowsResult": { "PaginatedRowsResult": {
@ -6146,7 +6141,7 @@
"error" "error"
] ]
}, },
"SpanWithChildren": { "SpanWithStatus": {
"type": "object", "type": "object",
"properties": { "properties": {
"span_id": { "span_id": {
@ -6194,12 +6189,6 @@
] ]
} }
}, },
"children": {
"type": "array",
"items": {
"$ref": "#/components/schemas/SpanWithChildren"
}
},
"status": { "status": {
"$ref": "#/components/schemas/SpanStatus" "$ref": "#/components/schemas/SpanStatus"
} }
@ -6209,8 +6198,7 @@
"span_id", "span_id",
"trace_id", "trace_id",
"name", "name",
"start_time", "start_time"
"children"
] ]
}, },
"Checkpoint": { "Checkpoint": {
@ -6236,31 +6224,11 @@
], ],
"title": "Artifacts of a finetuning job." "title": "Artifacts of a finetuning job."
}, },
"PostTrainingJobLogStream": { "JobStatus": {
"type": "object",
"properties": {
"job_uuid": {
"type": "string"
},
"log_lines": {
"type": "array",
"items": {
"type": "string"
}
}
},
"additionalProperties": false,
"required": [
"job_uuid",
"log_lines"
],
"title": "Stream of logs from a finetuning job."
},
"PostTrainingJobStatus": {
"type": "string", "type": "string",
"enum": [ "enum": [
"running",
"completed", "completed",
"in_progress",
"failed", "failed",
"scheduled" "scheduled"
] ]
@ -6272,7 +6240,7 @@
"type": "string" "type": "string"
}, },
"status": { "status": {
"$ref": "#/components/schemas/PostTrainingJobStatus" "$ref": "#/components/schemas/JobStatus"
}, },
"scheduled_at": { "scheduled_at": {
"type": "string", "type": "string",
@ -6456,13 +6424,6 @@
"job_id" "job_id"
] ]
}, },
"JobStatus": {
"type": "string",
"enum": [
"completed",
"in_progress"
]
},
"ProviderInfo": { "ProviderInfo": {
"type": "object", "type": "object",
"properties": { "properties": {
@ -6796,39 +6757,89 @@
"gamma" "gamma"
] ]
}, },
"DataConfig": {
"type": "object",
"properties": {
"dataset_id": {
"type": "string"
},
"batch_size": {
"type": "integer"
},
"shuffle": {
"type": "boolean"
},
"validation_dataset_id": {
"type": "string"
},
"packed": {
"type": "boolean",
"default": false
},
"train_on_input": {
"type": "boolean",
"default": false
}
},
"additionalProperties": false,
"required": [
"dataset_id",
"batch_size",
"shuffle"
]
},
"EfficiencyConfig": {
"type": "object",
"properties": {
"enable_activation_checkpointing": {
"type": "boolean",
"default": false
},
"enable_activation_offloading": {
"type": "boolean",
"default": false
},
"memory_efficient_fsdp_wrap": {
"type": "boolean",
"default": false
},
"fsdp_cpu_offload": {
"type": "boolean",
"default": false
}
},
"additionalProperties": false
},
"OptimizerConfig": { "OptimizerConfig": {
"type": "object", "type": "object",
"properties": { "properties": {
"optimizer_type": { "optimizer_type": {
"type": "string", "$ref": "#/components/schemas/OptimizerType"
"enum": [
"adam",
"adamw",
"sgd"
]
}, },
"lr": { "lr": {
"type": "number" "type": "number"
}, },
"lr_min": {
"type": "number"
},
"weight_decay": { "weight_decay": {
"type": "number" "type": "number"
},
"num_warmup_steps": {
"type": "integer"
} }
}, },
"additionalProperties": false, "additionalProperties": false,
"required": [ "required": [
"optimizer_type", "optimizer_type",
"lr", "lr",
"lr_min", "weight_decay",
"weight_decay" "num_warmup_steps"
] ]
}, },
"RLHFAlgorithm": { "OptimizerType": {
"type": "string", "type": "string",
"enum": [ "enum": [
"dpo" "adam",
"adamw",
"sgd"
] ]
}, },
"TrainingConfig": { "TrainingConfig": {
@ -6837,34 +6848,33 @@
"n_epochs": { "n_epochs": {
"type": "integer" "type": "integer"
}, },
"batch_size": { "max_steps_per_epoch": {
"type": "integer" "type": "integer"
}, },
"shuffle": { "gradient_accumulation_steps": {
"type": "boolean"
},
"n_iters": {
"type": "integer" "type": "integer"
}, },
"enable_activation_checkpointing": { "data_config": {
"type": "boolean" "$ref": "#/components/schemas/DataConfig"
}, },
"memory_efficient_fsdp_wrap": { "optimizer_config": {
"type": "boolean" "$ref": "#/components/schemas/OptimizerConfig"
}, },
"fsdp_cpu_offload": { "efficiency_config": {
"type": "boolean" "$ref": "#/components/schemas/EfficiencyConfig"
},
"dtype": {
"type": "string",
"default": "bf16"
} }
}, },
"additionalProperties": false, "additionalProperties": false,
"required": [ "required": [
"n_epochs", "n_epochs",
"batch_size", "max_steps_per_epoch",
"shuffle", "gradient_accumulation_steps",
"n_iters", "data_config",
"enable_activation_checkpointing", "optimizer_config"
"memory_efficient_fsdp_wrap",
"fsdp_cpu_offload"
] ]
}, },
"PreferenceOptimizeRequest": { "PreferenceOptimizeRequest": {
@ -6874,23 +6884,11 @@
"type": "string" "type": "string"
}, },
"finetuned_model": { "finetuned_model": {
"$ref": "#/components/schemas/URL"
},
"dataset_id": {
"type": "string" "type": "string"
}, },
"validation_dataset_id": {
"type": "string"
},
"algorithm": {
"$ref": "#/components/schemas/RLHFAlgorithm"
},
"algorithm_config": { "algorithm_config": {
"$ref": "#/components/schemas/DPOAlignmentConfig" "$ref": "#/components/schemas/DPOAlignmentConfig"
}, },
"optimizer_config": {
"$ref": "#/components/schemas/OptimizerConfig"
},
"training_config": { "training_config": {
"$ref": "#/components/schemas/TrainingConfig" "$ref": "#/components/schemas/TrainingConfig"
}, },
@ -6949,11 +6947,7 @@
"required": [ "required": [
"job_uuid", "job_uuid",
"finetuned_model", "finetuned_model",
"dataset_id",
"validation_dataset_id",
"algorithm",
"algorithm_config", "algorithm_config",
"optimizer_config",
"training_config", "training_config",
"hyperparam_search_config", "hyperparam_search_config",
"logger_config" "logger_config"
@ -7645,6 +7639,9 @@
} }
] ]
} }
},
"model_type": {
"$ref": "#/components/schemas/ModelType"
} }
}, },
"additionalProperties": false, "additionalProperties": false,
@ -8140,49 +8137,14 @@
"results" "results"
] ]
}, },
"DoraFinetuningConfig": {
"type": "object",
"properties": {
"lora_attn_modules": {
"type": "array",
"items": {
"type": "string"
}
},
"apply_lora_to_mlp": {
"type": "boolean"
},
"apply_lora_to_output": {
"type": "boolean"
},
"rank": {
"type": "integer"
},
"alpha": {
"type": "integer"
}
},
"additionalProperties": false,
"required": [
"lora_attn_modules",
"apply_lora_to_mlp",
"apply_lora_to_output",
"rank",
"alpha"
]
},
"FinetuningAlgorithm": {
"type": "string",
"enum": [
"full",
"lora",
"qlora",
"dora"
]
},
"LoraFinetuningConfig": { "LoraFinetuningConfig": {
"type": "object", "type": "object",
"properties": { "properties": {
"type": {
"type": "string",
"const": "LoRA",
"default": "LoRA"
},
"lora_attn_modules": { "lora_attn_modules": {
"type": "array", "type": "array",
"items": { "items": {
@ -8200,10 +8162,19 @@
}, },
"alpha": { "alpha": {
"type": "integer" "type": "integer"
},
"use_dora": {
"type": "boolean",
"default": false
},
"quantize_base": {
"type": "boolean",
"default": false
} }
}, },
"additionalProperties": false, "additionalProperties": false,
"required": [ "required": [
"type",
"lora_attn_modules", "lora_attn_modules",
"apply_lora_to_mlp", "apply_lora_to_mlp",
"apply_lora_to_output", "apply_lora_to_output",
@ -8211,35 +8182,26 @@
"alpha" "alpha"
] ]
}, },
"QLoraFinetuningConfig": { "QATFinetuningConfig": {
"type": "object", "type": "object",
"properties": { "properties": {
"lora_attn_modules": { "type": {
"type": "array", "type": "string",
"items": { "const": "QAT",
"default": "QAT"
},
"quantizer_name": {
"type": "string" "type": "string"
}
}, },
"apply_lora_to_mlp": { "group_size": {
"type": "boolean"
},
"apply_lora_to_output": {
"type": "boolean"
},
"rank": {
"type": "integer"
},
"alpha": {
"type": "integer" "type": "integer"
} }
}, },
"additionalProperties": false, "additionalProperties": false,
"required": [ "required": [
"lora_attn_modules", "type",
"apply_lora_to_mlp", "quantizer_name",
"apply_lora_to_output", "group_size"
"rank",
"alpha"
] ]
}, },
"SupervisedFineTuneRequest": { "SupervisedFineTuneRequest": {
@ -8248,34 +8210,6 @@
"job_uuid": { "job_uuid": {
"type": "string" "type": "string"
}, },
"model": {
"type": "string"
},
"dataset_id": {
"type": "string"
},
"validation_dataset_id": {
"type": "string"
},
"algorithm": {
"$ref": "#/components/schemas/FinetuningAlgorithm"
},
"algorithm_config": {
"oneOf": [
{
"$ref": "#/components/schemas/LoraFinetuningConfig"
},
{
"$ref": "#/components/schemas/QLoraFinetuningConfig"
},
{
"$ref": "#/components/schemas/DoraFinetuningConfig"
}
]
},
"optimizer_config": {
"$ref": "#/components/schemas/OptimizerConfig"
},
"training_config": { "training_config": {
"$ref": "#/components/schemas/TrainingConfig" "$ref": "#/components/schemas/TrainingConfig"
}, },
@ -8328,20 +8262,31 @@
} }
] ]
} }
},
"model": {
"type": "string"
},
"checkpoint_dir": {
"type": "string"
},
"algorithm_config": {
"oneOf": [
{
"$ref": "#/components/schemas/LoraFinetuningConfig"
},
{
"$ref": "#/components/schemas/QATFinetuningConfig"
}
]
} }
}, },
"additionalProperties": false, "additionalProperties": false,
"required": [ "required": [
"job_uuid", "job_uuid",
"model",
"dataset_id",
"validation_dataset_id",
"algorithm",
"algorithm_config",
"optimizer_config",
"training_config", "training_config",
"hyperparam_search_config", "hyperparam_search_config",
"logger_config" "logger_config",
"model"
] ]
}, },
"SyntheticDataGenerateRequest": { "SyntheticDataGenerateRequest": {
@ -8658,6 +8603,10 @@
"name": "DPOAlignmentConfig", "name": "DPOAlignmentConfig",
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/DPOAlignmentConfig\" />" "description": "<SchemaDefinition schemaRef=\"#/components/schemas/DPOAlignmentConfig\" />"
}, },
{
"name": "DataConfig",
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/DataConfig\" />"
},
{ {
"name": "Dataset", "name": "Dataset",
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/Dataset\" />" "description": "<SchemaDefinition schemaRef=\"#/components/schemas/Dataset\" />"
@ -8677,8 +8626,8 @@
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/DeleteAgentsSessionRequest\" />" "description": "<SchemaDefinition schemaRef=\"#/components/schemas/DeleteAgentsSessionRequest\" />"
}, },
{ {
"name": "DoraFinetuningConfig", "name": "EfficiencyConfig",
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/DoraFinetuningConfig\" />" "description": "<SchemaDefinition schemaRef=\"#/components/schemas/EfficiencyConfig\" />"
}, },
{ {
"name": "EmbeddingsRequest", "name": "EmbeddingsRequest",
@ -8706,10 +8655,6 @@
"name": "EvaluateRowsRequest", "name": "EvaluateRowsRequest",
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/EvaluateRowsRequest\" />" "description": "<SchemaDefinition schemaRef=\"#/components/schemas/EvaluateRowsRequest\" />"
}, },
{
"name": "FinetuningAlgorithm",
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/FinetuningAlgorithm\" />"
},
{ {
"name": "FunctionCallToolDefinition", "name": "FunctionCallToolDefinition",
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/FunctionCallToolDefinition\" />" "description": "<SchemaDefinition schemaRef=\"#/components/schemas/FunctionCallToolDefinition\" />"
@ -8826,6 +8771,10 @@
"name": "ModelCandidate", "name": "ModelCandidate",
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/ModelCandidate\" />" "description": "<SchemaDefinition schemaRef=\"#/components/schemas/ModelCandidate\" />"
}, },
{
"name": "ModelType",
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/ModelType\" />"
},
{ {
"name": "Models" "name": "Models"
}, },
@ -8833,6 +8782,10 @@
"name": "OptimizerConfig", "name": "OptimizerConfig",
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/OptimizerConfig\" />" "description": "<SchemaDefinition schemaRef=\"#/components/schemas/OptimizerConfig\" />"
}, },
{
"name": "OptimizerType",
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/OptimizerType\" />"
},
{ {
"name": "PaginatedRowsResult", "name": "PaginatedRowsResult",
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/PaginatedRowsResult\" />" "description": "<SchemaDefinition schemaRef=\"#/components/schemas/PaginatedRowsResult\" />"
@ -8852,14 +8805,6 @@
"name": "PostTrainingJobArtifactsResponse", "name": "PostTrainingJobArtifactsResponse",
"description": "Artifacts of a finetuning job.\n\n<SchemaDefinition schemaRef=\"#/components/schemas/PostTrainingJobArtifactsResponse\" />" "description": "Artifacts of a finetuning job.\n\n<SchemaDefinition schemaRef=\"#/components/schemas/PostTrainingJobArtifactsResponse\" />"
}, },
{
"name": "PostTrainingJobLogStream",
"description": "Stream of logs from a finetuning job.\n\n<SchemaDefinition schemaRef=\"#/components/schemas/PostTrainingJobLogStream\" />"
},
{
"name": "PostTrainingJobStatus",
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/PostTrainingJobStatus\" />"
},
{ {
"name": "PostTrainingJobStatusResponse", "name": "PostTrainingJobStatusResponse",
"description": "Status of a finetuning job.\n\n<SchemaDefinition schemaRef=\"#/components/schemas/PostTrainingJobStatusResponse\" />" "description": "Status of a finetuning job.\n\n<SchemaDefinition schemaRef=\"#/components/schemas/PostTrainingJobStatusResponse\" />"
@ -8873,8 +8818,8 @@
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/ProviderInfo\" />" "description": "<SchemaDefinition schemaRef=\"#/components/schemas/ProviderInfo\" />"
}, },
{ {
"name": "QLoraFinetuningConfig", "name": "QATFinetuningConfig",
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/QLoraFinetuningConfig\" />" "description": "<SchemaDefinition schemaRef=\"#/components/schemas/QATFinetuningConfig\" />"
}, },
{ {
"name": "QueryCondition", "name": "QueryCondition",
@ -8900,10 +8845,6 @@
"name": "QueryTracesRequest", "name": "QueryTracesRequest",
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/QueryTracesRequest\" />" "description": "<SchemaDefinition schemaRef=\"#/components/schemas/QueryTracesRequest\" />"
}, },
{
"name": "RLHFAlgorithm",
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/RLHFAlgorithm\" />"
},
{ {
"name": "RegexParserScoringFnParams", "name": "RegexParserScoringFnParams",
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/RegexParserScoringFnParams\" />" "description": "<SchemaDefinition schemaRef=\"#/components/schemas/RegexParserScoringFnParams\" />"
@ -9041,8 +8982,8 @@
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/SpanStatus\" />" "description": "<SchemaDefinition schemaRef=\"#/components/schemas/SpanStatus\" />"
}, },
{ {
"name": "SpanWithChildren", "name": "SpanWithStatus",
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/SpanWithChildren\" />" "description": "<SchemaDefinition schemaRef=\"#/components/schemas/SpanWithStatus\" />"
}, },
{ {
"name": "StopReason", "name": "StopReason",
@ -9237,16 +9178,16 @@
"CreateAgentSessionRequest", "CreateAgentSessionRequest",
"CreateAgentTurnRequest", "CreateAgentTurnRequest",
"DPOAlignmentConfig", "DPOAlignmentConfig",
"DataConfig",
"Dataset", "Dataset",
"DeleteAgentsRequest", "DeleteAgentsRequest",
"DeleteAgentsSessionRequest", "DeleteAgentsSessionRequest",
"DoraFinetuningConfig", "EfficiencyConfig",
"EmbeddingsRequest", "EmbeddingsRequest",
"EmbeddingsResponse", "EmbeddingsResponse",
"EvalTask", "EvalTask",
"EvaluateResponse", "EvaluateResponse",
"EvaluateRowsRequest", "EvaluateRowsRequest",
"FinetuningAlgorithm",
"FunctionCallToolDefinition", "FunctionCallToolDefinition",
"GetAgentsSessionRequest", "GetAgentsSessionRequest",
"GetSpanTreeRequest", "GetSpanTreeRequest",
@ -9273,24 +9214,23 @@
"MetricEvent", "MetricEvent",
"Model", "Model",
"ModelCandidate", "ModelCandidate",
"ModelType",
"OptimizerConfig", "OptimizerConfig",
"OptimizerType",
"PaginatedRowsResult", "PaginatedRowsResult",
"PhotogenToolDefinition", "PhotogenToolDefinition",
"PostTrainingJob", "PostTrainingJob",
"PostTrainingJobArtifactsResponse", "PostTrainingJobArtifactsResponse",
"PostTrainingJobLogStream",
"PostTrainingJobStatus",
"PostTrainingJobStatusResponse", "PostTrainingJobStatusResponse",
"PreferenceOptimizeRequest", "PreferenceOptimizeRequest",
"ProviderInfo", "ProviderInfo",
"QLoraFinetuningConfig", "QATFinetuningConfig",
"QueryCondition", "QueryCondition",
"QueryConditionOp", "QueryConditionOp",
"QueryDocumentsRequest", "QueryDocumentsRequest",
"QueryDocumentsResponse", "QueryDocumentsResponse",
"QuerySpansRequest", "QuerySpansRequest",
"QueryTracesRequest", "QueryTracesRequest",
"RLHFAlgorithm",
"RegexParserScoringFnParams", "RegexParserScoringFnParams",
"RegisterDatasetRequest", "RegisterDatasetRequest",
"RegisterEvalTaskRequest", "RegisterEvalTaskRequest",
@ -9322,7 +9262,7 @@
"SpanEndPayload", "SpanEndPayload",
"SpanStartPayload", "SpanStartPayload",
"SpanStatus", "SpanStatus",
"SpanWithChildren", "SpanWithStatus",
"StopReason", "StopReason",
"StructuredLogEvent", "StructuredLogEvent",
"SupervisedFineTuneRequest", "SupervisedFineTuneRequest",

308
docs/resources/llama-stack-spec.yaml
View file

@ -761,6 +761,28 @@ components:
- epsilon - epsilon
- gamma - gamma
type: object type: object
DataConfig:
additionalProperties: false
properties:
batch_size:
type: integer
dataset_id:
type: string
packed:
default: false
type: boolean
shuffle:
type: boolean
train_on_input:
default: false
type: boolean
validation_dataset_id:
type: string
required:
- dataset_id
- batch_size
- shuffle
type: object
Dataset: Dataset:
additionalProperties: false additionalProperties: false
properties: properties:
@ -908,27 +930,21 @@ components:
- agent_id - agent_id
- session_id - session_id
type: object type: object
DoraFinetuningConfig: EfficiencyConfig:
additionalProperties: false additionalProperties: false
properties: properties:
alpha: enable_activation_checkpointing:
type: integer default: false
apply_lora_to_mlp:
type: boolean type: boolean
apply_lora_to_output: enable_activation_offloading:
default: false
type: boolean
fsdp_cpu_offload:
default: false
type: boolean
memory_efficient_fsdp_wrap:
default: false
type: boolean type: boolean
lora_attn_modules:
items:
type: string
type: array
rank:
type: integer
required:
- lora_attn_modules
- apply_lora_to_mlp
- apply_lora_to_output
- rank
- alpha
type: object type: object
EmbeddingsRequest: EmbeddingsRequest:
additionalProperties: false additionalProperties: false
@ -1054,13 +1070,6 @@ components:
- scoring_functions - scoring_functions
- task_config - task_config
type: object type: object
FinetuningAlgorithm:
enum:
- full
- lora
- qlora
- dora
type: string
FunctionCallToolDefinition: FunctionCallToolDefinition:
additionalProperties: false additionalProperties: false
properties: properties:
@ -1230,6 +1239,8 @@ components:
enum: enum:
- completed - completed
- in_progress - in_progress
- failed
- scheduled
type: string type: string
KeyValueMemoryBank: KeyValueMemoryBank:
additionalProperties: false additionalProperties: false
@ -1358,9 +1369,20 @@ components:
items: items:
type: string type: string
type: array type: array
quantize_base:
default: false
type: boolean
rank: rank:
type: integer type: integer
type:
const: LoRA
default: LoRA
type: string
use_dora:
default: false
type: boolean
required: required:
- type
- lora_attn_modules - lora_attn_modules
- apply_lora_to_mlp - apply_lora_to_mlp
- apply_lora_to_output - apply_lora_to_output
@ -1621,6 +1643,9 @@ components:
- type: array - type: array
- type: object - type: object
type: object type: object
model_type:
$ref: '#/components/schemas/ModelType'
default: llm
provider_id: provider_id:
type: string type: string
provider_resource_id: provider_resource_id:
@ -1635,6 +1660,7 @@ components:
- provider_id - provider_id
- type - type
- metadata - metadata
- model_type
type: object type: object
ModelCandidate: ModelCandidate:
additionalProperties: false additionalProperties: false
@ -1654,27 +1680,34 @@ components:
- model - model
- sampling_params - sampling_params
type: object type: object
ModelType:
enum:
- llm
- embedding
type: string
OptimizerConfig: OptimizerConfig:
additionalProperties: false additionalProperties: false
properties: properties:
lr: lr:
type: number type: number
lr_min: num_warmup_steps:
type: number type: integer
optimizer_type: optimizer_type:
enum: $ref: '#/components/schemas/OptimizerType'
- adam
- adamw
- sgd
type: string
weight_decay: weight_decay:
type: number type: number
required: required:
- optimizer_type - optimizer_type
- lr - lr
- lr_min
- weight_decay - weight_decay
- num_warmup_steps
type: object type: object
OptimizerType:
enum:
- adam
- adamw
- sgd
type: string
PaginatedRowsResult: PaginatedRowsResult:
additionalProperties: false additionalProperties: false
properties: properties:
@ -1740,27 +1773,6 @@ components:
- checkpoints - checkpoints
title: Artifacts of a finetuning job. title: Artifacts of a finetuning job.
type: object type: object
PostTrainingJobLogStream:
additionalProperties: false
properties:
job_uuid:
type: string
log_lines:
items:
type: string
type: array
required:
- job_uuid
- log_lines
title: Stream of logs from a finetuning job.
type: object
PostTrainingJobStatus:
enum:
- running
- completed
- failed
- scheduled
type: string
PostTrainingJobStatusResponse: PostTrainingJobStatusResponse:
additionalProperties: false additionalProperties: false
properties: properties:
@ -1790,7 +1802,7 @@ components:
format: date-time format: date-time
type: string type: string
status: status:
$ref: '#/components/schemas/PostTrainingJobStatus' $ref: '#/components/schemas/JobStatus'
required: required:
- job_uuid - job_uuid
- status - status
@ -1800,14 +1812,10 @@ components:
PreferenceOptimizeRequest: PreferenceOptimizeRequest:
additionalProperties: false additionalProperties: false
properties: properties:
algorithm:
$ref: '#/components/schemas/RLHFAlgorithm'
algorithm_config: algorithm_config:
$ref: '#/components/schemas/DPOAlignmentConfig' $ref: '#/components/schemas/DPOAlignmentConfig'
dataset_id:
type: string
finetuned_model: finetuned_model:
$ref: '#/components/schemas/URL' type: string
hyperparam_search_config: hyperparam_search_config:
additionalProperties: additionalProperties:
oneOf: oneOf:
@ -1830,20 +1838,12 @@ components:
- type: array - type: array
- type: object - type: object
type: object type: object
optimizer_config:
$ref: '#/components/schemas/OptimizerConfig'
training_config: training_config:
$ref: '#/components/schemas/TrainingConfig' $ref: '#/components/schemas/TrainingConfig'
validation_dataset_id:
type: string
required: required:
- job_uuid - job_uuid
- finetuned_model - finetuned_model
- dataset_id
- validation_dataset_id
- algorithm
- algorithm_config - algorithm_config
- optimizer_config
- training_config - training_config
- hyperparam_search_config - hyperparam_search_config
- logger_config - logger_config
@ -1859,27 +1859,21 @@ components:
- provider_id - provider_id
- provider_type - provider_type
type: object type: object
QLoraFinetuningConfig: QATFinetuningConfig:
additionalProperties: false additionalProperties: false
properties: properties:
alpha: group_size:
type: integer type: integer
apply_lora_to_mlp: quantizer_name:
type: boolean type: string
apply_lora_to_output: type:
type: boolean const: QAT
lora_attn_modules: default: QAT
items:
type: string type: string
type: array
rank:
type: integer
required: required:
- lora_attn_modules - type
- apply_lora_to_mlp - quantizer_name
- apply_lora_to_output - group_size
- rank
- alpha
type: object type: object
QueryCondition: QueryCondition:
additionalProperties: false additionalProperties: false
@ -2003,10 +1997,6 @@ components:
type: string type: string
type: array type: array
type: object type: object
RLHFAlgorithm:
enum:
- dpo
type: string
RegexParserScoringFnParams: RegexParserScoringFnParams:
additionalProperties: false additionalProperties: false
properties: properties:
@ -2209,6 +2199,8 @@ components:
type: object type: object
model_id: model_id:
type: string type: string
model_type:
$ref: '#/components/schemas/ModelType'
provider_id: provider_id:
type: string type: string
provider_model_id: provider_model_id:
@ -2941,7 +2933,7 @@ components:
- ok - ok
- error - error
type: string type: string
SpanWithChildren: SpanWithStatus:
additionalProperties: false additionalProperties: false
properties: properties:
attributes: attributes:
@ -2954,10 +2946,6 @@ components:
- type: array - type: array
- type: object - type: object
type: object type: object
children:
items:
$ref: '#/components/schemas/SpanWithChildren'
type: array
end_time: end_time:
format: date-time format: date-time
type: string type: string
@ -2979,7 +2967,6 @@ components:
- trace_id - trace_id
- name - name
- start_time - start_time
- children
type: object type: object
StopReason: StopReason:
enum: enum:
@ -3025,14 +3012,11 @@ components:
SupervisedFineTuneRequest: SupervisedFineTuneRequest:
additionalProperties: false additionalProperties: false
properties: properties:
algorithm:
$ref: '#/components/schemas/FinetuningAlgorithm'
algorithm_config: algorithm_config:
oneOf: oneOf:
- $ref: '#/components/schemas/LoraFinetuningConfig' - $ref: '#/components/schemas/LoraFinetuningConfig'
- $ref: '#/components/schemas/QLoraFinetuningConfig' - $ref: '#/components/schemas/QATFinetuningConfig'
- $ref: '#/components/schemas/DoraFinetuningConfig' checkpoint_dir:
dataset_id:
type: string type: string
hyperparam_search_config: hyperparam_search_config:
additionalProperties: additionalProperties:
@ -3058,23 +3042,14 @@ components:
type: object type: object
model: model:
type: string type: string
optimizer_config:
$ref: '#/components/schemas/OptimizerConfig'
training_config: training_config:
$ref: '#/components/schemas/TrainingConfig' $ref: '#/components/schemas/TrainingConfig'
validation_dataset_id:
type: string
required: required:
- job_uuid - job_uuid
- model
- dataset_id
- validation_dataset_id
- algorithm
- algorithm_config
- optimizer_config
- training_config - training_config
- hyperparam_search_config - hyperparam_search_config
- logger_config - logger_config
- model
type: object type: object
SyntheticDataGenerateRequest: SyntheticDataGenerateRequest:
additionalProperties: false additionalProperties: false
@ -3384,28 +3359,27 @@ components:
TrainingConfig: TrainingConfig:
additionalProperties: false additionalProperties: false
properties: properties:
batch_size: data_config:
$ref: '#/components/schemas/DataConfig'
dtype:
default: bf16
type: string
efficiency_config:
$ref: '#/components/schemas/EfficiencyConfig'
gradient_accumulation_steps:
type: integer
max_steps_per_epoch:
type: integer type: integer
enable_activation_checkpointing:
type: boolean
fsdp_cpu_offload:
type: boolean
memory_efficient_fsdp_wrap:
type: boolean
n_epochs: n_epochs:
type: integer type: integer
n_iters: optimizer_config:
type: integer $ref: '#/components/schemas/OptimizerConfig'
shuffle:
type: boolean
required: required:
- n_epochs - n_epochs
- batch_size - max_steps_per_epoch
- shuffle - gradient_accumulation_steps
- n_iters - data_config
- enable_activation_checkpointing - optimizer_config
- memory_efficient_fsdp_wrap
- fsdp_cpu_offload
type: object type: object
Turn: Turn:
additionalProperties: false additionalProperties: false
@ -3548,6 +3522,9 @@ components:
properties: properties:
chunk_size_in_tokens: chunk_size_in_tokens:
type: integer type: integer
embedding_dimension:
default: 384
type: integer
embedding_model: embedding_model:
type: string type: string
identifier: identifier:
@ -4601,7 +4578,9 @@ paths:
content: content:
application/json: application/json:
schema: schema:
$ref: '#/components/schemas/PostTrainingJobArtifactsResponse' oneOf:
- $ref: '#/components/schemas/PostTrainingJobArtifactsResponse'
- type: 'null'
description: OK description: OK
tags: tags:
- PostTraining (Coming Soon) - PostTraining (Coming Soon)
@ -4626,30 +4605,6 @@ paths:
description: OK description: OK
tags: tags:
- PostTraining (Coming Soon) - PostTraining (Coming Soon)
/alpha/post-training/job/logs:
get:
parameters:
- in: query
name: job_uuid
required: true
schema:
type: string
- description: JSON-encoded provider data which will be made available to the
adapter servicing the API
in: header
name: X-LlamaStack-ProviderData
required: false
schema:
type: string
responses:
'200':
content:
application/json:
schema:
$ref: '#/components/schemas/PostTrainingJobLogStream'
description: OK
tags:
- PostTraining (Coming Soon)
/alpha/post-training/job/status: /alpha/post-training/job/status:
get: get:
parameters: parameters:
@ -4670,7 +4625,9 @@ paths:
content: content:
application/json: application/json:
schema: schema:
$ref: '#/components/schemas/PostTrainingJobStatusResponse' oneOf:
- $ref: '#/components/schemas/PostTrainingJobStatusResponse'
- type: 'null'
description: OK description: OK
tags: tags:
- PostTraining (Coming Soon) - PostTraining (Coming Soon)
@ -5054,7 +5011,9 @@ paths:
content: content:
application/json: application/json:
schema: schema:
$ref: '#/components/schemas/SpanWithChildren' additionalProperties:
$ref: '#/components/schemas/SpanWithStatus'
type: object
description: OK description: OK
tags: tags:
- Telemetry - Telemetry
@ -5290,6 +5249,8 @@ tags:
- description: <SchemaDefinition schemaRef="#/components/schemas/DPOAlignmentConfig" - description: <SchemaDefinition schemaRef="#/components/schemas/DPOAlignmentConfig"
/> />
name: DPOAlignmentConfig name: DPOAlignmentConfig
- description: <SchemaDefinition schemaRef="#/components/schemas/DataConfig" />
name: DataConfig
- description: <SchemaDefinition schemaRef="#/components/schemas/Dataset" /> - description: <SchemaDefinition schemaRef="#/components/schemas/Dataset" />
name: Dataset name: Dataset
- name: DatasetIO - name: DatasetIO
@ -5300,9 +5261,9 @@ tags:
- description: <SchemaDefinition schemaRef="#/components/schemas/DeleteAgentsSessionRequest" - description: <SchemaDefinition schemaRef="#/components/schemas/DeleteAgentsSessionRequest"
/> />
name: DeleteAgentsSessionRequest name: DeleteAgentsSessionRequest
- description: <SchemaDefinition schemaRef="#/components/schemas/DoraFinetuningConfig" - description: <SchemaDefinition schemaRef="#/components/schemas/EfficiencyConfig"
/> />
name: DoraFinetuningConfig name: EfficiencyConfig
- description: <SchemaDefinition schemaRef="#/components/schemas/EmbeddingsRequest" - description: <SchemaDefinition schemaRef="#/components/schemas/EmbeddingsRequest"
/> />
name: EmbeddingsRequest name: EmbeddingsRequest
@ -5319,9 +5280,6 @@ tags:
- description: <SchemaDefinition schemaRef="#/components/schemas/EvaluateRowsRequest" - description: <SchemaDefinition schemaRef="#/components/schemas/EvaluateRowsRequest"
/> />
name: EvaluateRowsRequest name: EvaluateRowsRequest
- description: <SchemaDefinition schemaRef="#/components/schemas/FinetuningAlgorithm"
/>
name: FinetuningAlgorithm
- description: <SchemaDefinition schemaRef="#/components/schemas/FunctionCallToolDefinition" - description: <SchemaDefinition schemaRef="#/components/schemas/FunctionCallToolDefinition"
/> />
name: FunctionCallToolDefinition name: FunctionCallToolDefinition
@ -5395,10 +5353,14 @@ tags:
name: Model name: Model
- description: <SchemaDefinition schemaRef="#/components/schemas/ModelCandidate" /> - description: <SchemaDefinition schemaRef="#/components/schemas/ModelCandidate" />
name: ModelCandidate name: ModelCandidate
- description: <SchemaDefinition schemaRef="#/components/schemas/ModelType" />
name: ModelType
- name: Models - name: Models
- description: <SchemaDefinition schemaRef="#/components/schemas/OptimizerConfig" - description: <SchemaDefinition schemaRef="#/components/schemas/OptimizerConfig"
/> />
name: OptimizerConfig name: OptimizerConfig
- description: <SchemaDefinition schemaRef="#/components/schemas/OptimizerType" />
name: OptimizerType
- description: <SchemaDefinition schemaRef="#/components/schemas/PaginatedRowsResult" - description: <SchemaDefinition schemaRef="#/components/schemas/PaginatedRowsResult"
/> />
name: PaginatedRowsResult name: PaginatedRowsResult
@ -5415,14 +5377,6 @@ tags:
<SchemaDefinition schemaRef="#/components/schemas/PostTrainingJobArtifactsResponse" <SchemaDefinition schemaRef="#/components/schemas/PostTrainingJobArtifactsResponse"
/>' />'
name: PostTrainingJobArtifactsResponse name: PostTrainingJobArtifactsResponse
- description: 'Stream of logs from a finetuning job.
<SchemaDefinition schemaRef="#/components/schemas/PostTrainingJobLogStream" />'
name: PostTrainingJobLogStream
- description: <SchemaDefinition schemaRef="#/components/schemas/PostTrainingJobStatus"
/>
name: PostTrainingJobStatus
- description: 'Status of a finetuning job. - description: 'Status of a finetuning job.
@ -5434,9 +5388,9 @@ tags:
name: PreferenceOptimizeRequest name: PreferenceOptimizeRequest
- description: <SchemaDefinition schemaRef="#/components/schemas/ProviderInfo" /> - description: <SchemaDefinition schemaRef="#/components/schemas/ProviderInfo" />
name: ProviderInfo name: ProviderInfo
- description: <SchemaDefinition schemaRef="#/components/schemas/QLoraFinetuningConfig" - description: <SchemaDefinition schemaRef="#/components/schemas/QATFinetuningConfig"
/> />
name: QLoraFinetuningConfig name: QATFinetuningConfig
- description: <SchemaDefinition schemaRef="#/components/schemas/QueryCondition" /> - description: <SchemaDefinition schemaRef="#/components/schemas/QueryCondition" />
name: QueryCondition name: QueryCondition
- description: <SchemaDefinition schemaRef="#/components/schemas/QueryConditionOp" - description: <SchemaDefinition schemaRef="#/components/schemas/QueryConditionOp"
@ -5454,8 +5408,6 @@ tags:
- description: <SchemaDefinition schemaRef="#/components/schemas/QueryTracesRequest" - description: <SchemaDefinition schemaRef="#/components/schemas/QueryTracesRequest"
/> />
name: QueryTracesRequest name: QueryTracesRequest
- description: <SchemaDefinition schemaRef="#/components/schemas/RLHFAlgorithm" />
name: RLHFAlgorithm
- description: <SchemaDefinition schemaRef="#/components/schemas/RegexParserScoringFnParams" - description: <SchemaDefinition schemaRef="#/components/schemas/RegexParserScoringFnParams"
/> />
name: RegexParserScoringFnParams name: RegexParserScoringFnParams
@ -5545,9 +5497,8 @@ tags:
name: SpanStartPayload name: SpanStartPayload
- description: <SchemaDefinition schemaRef="#/components/schemas/SpanStatus" /> - description: <SchemaDefinition schemaRef="#/components/schemas/SpanStatus" />
name: SpanStatus name: SpanStatus
- description: <SchemaDefinition schemaRef="#/components/schemas/SpanWithChildren" - description: <SchemaDefinition schemaRef="#/components/schemas/SpanWithStatus" />
/> name: SpanWithStatus
name: SpanWithChildren
- description: <SchemaDefinition schemaRef="#/components/schemas/StopReason" /> - description: <SchemaDefinition schemaRef="#/components/schemas/StopReason" />
name: StopReason name: StopReason
- description: <SchemaDefinition schemaRef="#/components/schemas/StructuredLogEvent" - description: <SchemaDefinition schemaRef="#/components/schemas/StructuredLogEvent"
@ -5703,16 +5654,16 @@ x-tagGroups:
- CreateAgentSessionRequest - CreateAgentSessionRequest
- CreateAgentTurnRequest - CreateAgentTurnRequest
- DPOAlignmentConfig - DPOAlignmentConfig
- DataConfig
- Dataset - Dataset
- DeleteAgentsRequest - DeleteAgentsRequest
- DeleteAgentsSessionRequest - DeleteAgentsSessionRequest
- DoraFinetuningConfig - EfficiencyConfig
- EmbeddingsRequest - EmbeddingsRequest
- EmbeddingsResponse - EmbeddingsResponse
- EvalTask - EvalTask
- EvaluateResponse - EvaluateResponse
- EvaluateRowsRequest - EvaluateRowsRequest
- FinetuningAlgorithm
- FunctionCallToolDefinition - FunctionCallToolDefinition
- GetAgentsSessionRequest - GetAgentsSessionRequest
- GetSpanTreeRequest - GetSpanTreeRequest
@ -5739,24 +5690,23 @@ x-tagGroups:
- MetricEvent - MetricEvent
- Model - Model
- ModelCandidate - ModelCandidate
- ModelType
- OptimizerConfig - OptimizerConfig
- OptimizerType
- PaginatedRowsResult - PaginatedRowsResult
- PhotogenToolDefinition - PhotogenToolDefinition
- PostTrainingJob - PostTrainingJob
- PostTrainingJobArtifactsResponse - PostTrainingJobArtifactsResponse
- PostTrainingJobLogStream
- PostTrainingJobStatus
- PostTrainingJobStatusResponse - PostTrainingJobStatusResponse
- PreferenceOptimizeRequest - PreferenceOptimizeRequest
- ProviderInfo - ProviderInfo
- QLoraFinetuningConfig - QATFinetuningConfig
- QueryCondition - QueryCondition
- QueryConditionOp - QueryConditionOp
- QueryDocumentsRequest - QueryDocumentsRequest
- QueryDocumentsResponse - QueryDocumentsResponse
- QuerySpansRequest - QuerySpansRequest
- QueryTracesRequest - QueryTracesRequest
- RLHFAlgorithm
- RegexParserScoringFnParams - RegexParserScoringFnParams
- RegisterDatasetRequest - RegisterDatasetRequest
- RegisterEvalTaskRequest - RegisterEvalTaskRequest
@ -5788,7 +5738,7 @@ x-tagGroups:
- SpanEndPayload - SpanEndPayload
- SpanStartPayload - SpanStartPayload
- SpanStatus - SpanStatus
- SpanWithChildren - SpanWithStatus
- StopReason - StopReason
- StructuredLogEvent - StructuredLogEvent
- SupervisedFineTuneRequest - SupervisedFineTuneRequest

5
docs/source/building_applications/index.md
View file

@ -404,8 +404,9 @@ traces = client.telemetry.query_traces(
}] }]
) )
# Get detailed span information # Get spans within the root span; indexed by ID
span_tree = client.telemetry.get_span_tree( # Use parent_span_id to build a tree out of it
spans_by_id = client.telemetry.get_span_tree(
span_id=traces[0].root_span_id span_id=traces[0].root_span_id
) )
``` ```

5
llama_stack/apis/telemetry/telemetry.py
View file

@ -150,8 +150,7 @@ class EvalTrace(BaseModel):
@json_schema_type @json_schema_type
class SpanWithChildren(Span): class SpanWithStatus(Span):
children: List["SpanWithChildren"] = Field(default_factory=list)
status: Optional[SpanStatus] = None status: Optional[SpanStatus] = None
@ -192,7 +191,7 @@ class Telemetry(Protocol):
span_id: str, span_id: str,
attributes_to_return: Optional[List[str]] = None, attributes_to_return: Optional[List[str]] = None,
max_depth: Optional[int] = None, max_depth: Optional[int] = None,
) -> SpanWithChildren: ... ) -> Dict[str, SpanWithStatus]: ...
@webmethod(route="/telemetry/query-spans", method="POST") @webmethod(route="/telemetry/query-spans", method="POST")
async def query_spans( async def query_spans(

2
llama_stack/providers/inline/telemetry/meta_reference/telemetry.py
View file

@ -243,7 +243,7 @@ class TelemetryAdapter(TelemetryDatasetMixin, Telemetry):
span_id: str, span_id: str,
attributes_to_return: Optional[List[str]] = None, attributes_to_return: Optional[List[str]] = None,
max_depth: Optional[int] = None, max_depth: Optional[int] = None,
) -> SpanWithChildren: ) -> Dict[str, SpanWithStatus]:
return await self.trace_store.get_span_tree( return await self.trace_store.get_span_tree(
span_id=span_id, span_id=span_id,
attributes_to_return=attributes_to_return, attributes_to_return=attributes_to_return,

16
llama_stack/providers/utils/telemetry/dataset_mixin.py
View file

@ -7,7 +7,7 @@
from typing import List, Optional from typing import List, Optional
from llama_stack.apis.datasetio import DatasetIO from llama_stack.apis.datasetio import DatasetIO
from llama_stack.apis.telemetry import QueryCondition, Span, SpanWithChildren from llama_stack.apis.telemetry import QueryCondition, Span
class TelemetryDatasetMixin: class TelemetryDatasetMixin:
@ -53,19 +53,18 @@ class TelemetryDatasetMixin:
spans = [] spans = []
for trace in traces: for trace in traces:
span_tree = await self.get_span_tree( spans_by_id = await self.get_span_tree(
span_id=trace.root_span_id, span_id=trace.root_span_id,
attributes_to_return=attributes_to_return, attributes_to_return=attributes_to_return,
max_depth=max_depth, max_depth=max_depth,
) )
def extract_spans(span: SpanWithChildren) -> List[Span]: for span in spans_by_id.values():
result = []
if span.attributes and all( if span.attributes and all(
attr in span.attributes and span.attributes[attr] is not None attr in span.attributes and span.attributes[attr] is not None
for attr in attributes_to_return for attr in attributes_to_return
): ):
result.append( spans.append(
Span( Span(
trace_id=trace.root_span_id, trace_id=trace.root_span_id,
span_id=span.span_id, span_id=span.span_id,
@ -77,11 +76,4 @@ class TelemetryDatasetMixin:
) )
) )
for child in span.children:
result.extend(extract_spans(child))
return result
spans.extend(extract_spans(span_tree))
return spans return spans

23
llama_stack/providers/utils/telemetry/sqlite_trace_store.py
View file

@ -6,11 +6,11 @@
import json import json
from datetime import datetime from datetime import datetime
from typing import List, Optional, Protocol from typing import Dict, List, Optional, Protocol
import aiosqlite import aiosqlite
from llama_stack.apis.telemetry import QueryCondition, SpanWithChildren, Trace from llama_stack.apis.telemetry import QueryCondition, SpanWithStatus, Trace
class TraceStore(Protocol): class TraceStore(Protocol):
@ -27,7 +27,7 @@ class TraceStore(Protocol):
span_id: str, span_id: str,
attributes_to_return: Optional[List[str]] = None, attributes_to_return: Optional[List[str]] = None,
max_depth: Optional[int] = None, max_depth: Optional[int] = None,
) -> SpanWithChildren: ... ) -> Dict[str, SpanWithStatus]: ...
class SQLiteTraceStore(TraceStore): class SQLiteTraceStore(TraceStore):
@ -114,7 +114,7 @@ class SQLiteTraceStore(TraceStore):
span_id: str, span_id: str,
attributes_to_return: Optional[List[str]] = None, attributes_to_return: Optional[List[str]] = None,
max_depth: Optional[int] = None, max_depth: Optional[int] = None,
) -> SpanWithChildren: ) -> Dict[str, SpanWithStatus]:
# Build the attributes selection # Build the attributes selection
attributes_select = "s.attributes" attributes_select = "s.attributes"
if attributes_to_return: if attributes_to_return:
@ -143,6 +143,7 @@ class SQLiteTraceStore(TraceStore):
ORDER BY depth, start_time ORDER BY depth, start_time
""" """
spans_by_id = {}
async with aiosqlite.connect(self.conn_string) as conn: async with aiosqlite.connect(self.conn_string) as conn:
conn.row_factory = aiosqlite.Row conn.row_factory = aiosqlite.Row
async with conn.execute(query, (span_id, max_depth, max_depth)) as cursor: async with conn.execute(query, (span_id, max_depth, max_depth)) as cursor:
@ -151,12 +152,8 @@ class SQLiteTraceStore(TraceStore):
if not rows: if not rows:
raise ValueError(f"Span {span_id} not found") raise ValueError(f"Span {span_id} not found")
# Build span tree
spans_by_id = {}
root_span = None
for row in rows: for row in rows:
span = SpanWithChildren( span = SpanWithStatus(
span_id=row["span_id"], span_id=row["span_id"],
trace_id=row["trace_id"], trace_id=row["trace_id"],
parent_span_id=row["parent_span_id"], parent_span_id=row["parent_span_id"],
@ -165,14 +162,8 @@ class SQLiteTraceStore(TraceStore):
end_time=datetime.fromisoformat(row["end_time"]), end_time=datetime.fromisoformat(row["end_time"]),
attributes=json.loads(row["filtered_attributes"]), attributes=json.loads(row["filtered_attributes"]),
status=row["status"].lower(), status=row["status"].lower(),
children=[],
) )
spans_by_id[span.span_id] = span spans_by_id[span.span_id] = span
if span.span_id == span_id: return spans_by_id
root_span = span
elif span.parent_span_id in spans_by_id:
spans_by_id[span.parent_span_id].children.append(span)
return root_span

8
llama_stack/providers/utils/telemetry/trace_protocol.py
View file

@ -41,8 +41,6 @@ def trace_protocol(cls: Type[T]) -> Type[T]:
""" """
def trace_method(method: Callable) -> Callable: def trace_method(method: Callable) -> Callable:
from llama_stack.providers.utils.telemetry import tracing
is_async = asyncio.iscoroutinefunction(method) is_async = asyncio.iscoroutinefunction(method)
is_async_gen = inspect.isasyncgenfunction(method) is_async_gen = inspect.isasyncgenfunction(method)
@ -77,6 +75,8 @@ def trace_protocol(cls: Type[T]) -> Type[T]:
async def async_gen_wrapper( async def async_gen_wrapper(
self: Any, *args: Any, **kwargs: Any self: Any, *args: Any, **kwargs: Any
) -> AsyncGenerator: ) -> AsyncGenerator:
from llama_stack.providers.utils.telemetry import tracing
class_name, method_name, span_attributes = create_span_context( class_name, method_name, span_attributes = create_span_context(
self, *args, **kwargs self, *args, **kwargs
) )
@ -92,6 +92,8 @@ def trace_protocol(cls: Type[T]) -> Type[T]:
@wraps(method) @wraps(method)
async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
from llama_stack.providers.utils.telemetry import tracing
class_name, method_name, span_attributes = create_span_context( class_name, method_name, span_attributes = create_span_context(
self, *args, **kwargs self, *args, **kwargs
) )
@ -107,6 +109,8 @@ def trace_protocol(cls: Type[T]) -> Type[T]:
@wraps(method) @wraps(method)
def sync_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: def sync_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
from llama_stack.providers.utils.telemetry import tracing
class_name, method_name, span_attributes = create_span_context( class_name, method_name, span_attributes = create_span_context(
self, *args, **kwargs self, *args, **kwargs
) )