Thanks to visit codestin.com
Credit goes to github.com

Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 20 additions & 1 deletion cli/exp_scaletest_chat.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"sync"
"time"

"github.com/google/uuid"
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/promhttp"
"golang.org/x/xerrors"
Expand Down Expand Up @@ -77,10 +78,23 @@ func (r *RootCmd) scaletestChat() *serpent.Command {
}

logger := inv.Logger
modelID, err := chat.EnsureScaletestChatModel(ctx, client, logger, llmMockURL, providerPropagationWait)
modelForOrg, err := chat.EnsureScaletestChatModel(ctx, client, logger, llmMockURL, providerPropagationWait)
if err != nil {
return err
}
uncachedModelForOrg := modelForOrg
modelIDsByOrganization := make(map[uuid.UUID]uuid.UUID)
modelForOrg = func(organizationID uuid.UUID) (uuid.UUID, error) {
if modelID, ok := modelIDsByOrganization[organizationID]; ok {
return modelID, nil
}
modelID, err := uncachedModelForOrg(organizationID)
if err != nil {
return uuid.Nil, err
}
modelIDsByOrganization[organizationID] = modelID
return modelID, nil
}

// Start metrics and tracing before creating runners.
reg := prometheus.NewRegistry()
Expand Down Expand Up @@ -124,6 +138,11 @@ func (r *RootCmd) scaletestChat() *serpent.Command {
turnStartReadyWaitGroup.Add(1)
}

modelID, err := modelForOrg(targetWorkspace.OrganizationID)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Cache scaletest model IDs per organization

When a scale test creates multiple chats for the same workspace or organization, this call runs once per chat inside both loops, and each invocation performs a synchronous ChatModels request before the harness starts. Large --chats-per-workspace runs consequently issue hundreds or thousands of identical serial bootstrap requests, substantially delaying or stalling the load test itself. Resolve each organization's model once and reuse the ID for all of its runners.

Useful? React with 👍 / 👎.

if err != nil {
return xerrors.Errorf("ensure scaletest model config for organization %s: %w", targetWorkspace.OrganizationID, err)
}

cfg := chat.Config{
OrganizationID: targetWorkspace.OrganizationID,
WorkspaceID: targetWorkspace.ID,
Expand Down
8 changes: 5 additions & 3 deletions cli/exp_scaletest_chat_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -77,9 +77,11 @@ func TestScaleTestChat(t *testing.T) {
require.Equal(t, mockURL, provider.BaseURL)

expClient := codersdk.NewExperimentalClient(client)
configs, err := expClient.ChatModels(ctx)
defaultOrg, err := client.OrganizationByName(ctx, codersdk.DefaultOrganization)
require.NoError(t, err)
matchingConfigs := scaletestModelsForProvider(configs, provider.ID)
configs, err := expClient.ChatModels(ctx, defaultOrg.ID)
require.NoError(t, err)
matchingConfigs := scaletestModelConfigsForProvider(configs.Models, provider.ID)
require.Len(t, matchingConfigs, 1)
require.True(t, matchingConfigs[0].Enabled)

Expand Down Expand Up @@ -125,7 +127,7 @@ func chatMessageText(messages []codersdk.ChatMessage, role codersdk.ChatMessageR
return b.String(), found
}

func scaletestModelsForProvider(configs []codersdk.ChatModel, providerID uuid.UUID) []codersdk.ChatModel {
func scaletestModelConfigsForProvider(configs []codersdk.ChatModel, providerID uuid.UUID) []codersdk.ChatModel {
matches := make([]codersdk.ChatModel, 0, 1)
for _, config := range configs {
if config.AIProviderID != providerID {
Expand Down
Loading
Loading