From bc0c76d79dccf1e8a2cd408ce5a2868712ff621f Mon Sep 17 00:00:00 2001 From: RUFFY-369 Date: Thu, 28 Nov 2024 21:16:28 +0530 Subject: [PATCH 1/7] chore:add support for JSON string to defaultImage for overriding pipeline specific images in mappings --- worker/docker.go | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/worker/docker.go b/worker/docker.go index 893d371c..d5031caa 100644 --- a/worker/docker.go +++ b/worker/docker.go @@ -99,6 +99,33 @@ type DockerManager struct { mu *sync.Mutex } +func overridePipelineImages(defaultImage string) error { + // First check if the defaultImage is a valid JSON string + var pipelineOverrides map[string]string + if strings.HasPrefix(defaultImage, "{") && strings.HasSuffix(defaultImage, "}") { + // Parse the JSON string for custom pipeline images override + if err := json.Unmarshal([]byte(defaultImage), &pipelineOverrides); err != nil { + return fmt.Errorf("failed to parse pipeline overrides JSON: %w", err) + } + + // Override the pipelineToImage and livePipelineToImage mappings + for pipeline, image := range pipelineOverrides { + if _, isPresent := pipelineToImage[pipeline]; isPresent { + pipelineToImage[pipeline] = image + } else { + slog.Warn("Pipeline not found in pipelineToImage map", "pipeline", pipeline) + } + + if _, isPresent := livePipelineToImage[pipeline]; isPresent { + livePipelineToImage[pipeline] = image + } else { + slog.Warn("Pipeline not found in livePipelineToImage map", "pipeline", pipeline) + } + } + } + return nil +} + func NewDockerManager(defaultImage string, gpus []string, modelDir string, client DockerClient) (*DockerManager, error) { ctx, cancel := context.WithTimeout(context.Background(), containerTimeout) if err := removeExistingContainers(ctx, client); err != nil { @@ -107,6 +134,11 @@ func NewDockerManager(defaultImage string, gpus []string, modelDir string, clien } cancel() + // call to handle image overriding logic + if err := overridePipelineImages(defaultImage); err != nil { + return nil, err + } + manager := &DockerManager{ defaultImage: defaultImage, gpus: gpus, From cbe8ec4c1ba1d9ed0521269142e899cec1da4321 Mon Sep 17 00:00:00 2001 From: RUFFY-369 Date: Sun, 1 Dec 2024 20:41:11 +0530 Subject: [PATCH 2/7] chore:add docstring for overridePipelineImages function --- worker/docker.go | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/worker/docker.go b/worker/docker.go index d5031caa..a00507b8 100644 --- a/worker/docker.go +++ b/worker/docker.go @@ -100,6 +100,14 @@ type DockerManager struct { } func overridePipelineImages(defaultImage string) error { + // overridePipelineImages function parses a JSON string containing pipeline-to-image mappings and overrides the default mappings if valid. + // It updates the `pipelineToImage` and `livePipelineToImage` maps with custom images. + // Parameters: + // - defaultImage: A string that can either be containerImage name or a JSON string with overrides for pipeline-to-image mappings. + // + // Returns: + // - error: An error if the JSON parsing fails or if the mapping is not found in existing maps else `nil`. + // First check if the defaultImage is a valid JSON string var pipelineOverrides map[string]string if strings.HasPrefix(defaultImage, "{") && strings.HasSuffix(defaultImage, "}") { From 952a58824ce265110ef46b6c43bf26982712e151 Mon Sep 17 00:00:00 2001 From: RUFFY-369 Date: Sun, 1 Dec 2024 21:28:11 +0530 Subject: [PATCH 3/7] test:add test for overridePipelineImages function --- worker/docker_test.go | 65 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/worker/docker_test.go b/worker/docker_test.go index cbb20086..8c405383 100644 --- a/worker/docker_test.go +++ b/worker/docker_test.go @@ -784,3 +784,68 @@ func TestDockerWaitUntilRunning(t *testing.T) { mockDockerClient.AssertExpectations(t) }) } + +func TestDockerManager_overridePipelineImages(t *testing.T) { + mockDockerClient := new(MockDockerClient) + dockerManager := createDockerManager(mockDockerClient) + + tests := []struct { + name string + inputJSON string + pipeline string + expectedImage string + expectError bool + }{ + { + name: "ValidOverride", + inputJSON: `{"segment-anything-2": "custom-image:1.0"}`, + pipeline: "segment-anything-2", + expectedImage: "custom-image:1.0", + expectError: false, + }, + { + name: "MultipleOverrides", + inputJSON: `{"segment-anything-2": "custom-image:1.0", "text-to-speech": "speech-image:2.0"}`, + pipeline: "text-to-speech", + expectedImage: "speech-image:2.0", + expectError: false, + }, + { + name: "NoOverrideFallback", + inputJSON: `{"segment-anything-2": "custom-image:1.0"}`, + pipeline: "video-to-video", + expectedImage: "default-image", + expectError: false, + }, + { + name: "EmptyJSON", + inputJSON: `{}`, + pipeline: "segment-anything-2", + expectedImage: "default-image", + expectError: false, + }, + { + name: "MalformedJSON", + inputJSON: `{"segment-anything-2": "custom-image:1.0`, + pipeline: "segment-anything-2", + expectError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Call overridePipelineImages function with the mock data. + err := overridePipelineImages(tt.inputJSON) + + if tt.expectError { + require.Error(t, err) + } else { + require.NoError(t, err) + + // Verify the expected image. + image, _ := dockerManager.getContainerImageName(tt.pipeline, "") + require.Equal(t, tt.expectedImage, image) + } + }) + } +} From ec397d855b6a98e2677a6669a4bcf35c8c8d88d9 Mon Sep 17 00:00:00 2001 From: RUFFY-369 Date: Mon, 2 Dec 2024 14:18:22 +0530 Subject: [PATCH 4/7] fix:malformed json test case handling;nits on expectedimage name --- worker/docker.go | 8 +++++--- worker/docker_test.go | 6 +++--- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/worker/docker.go b/worker/docker.go index a00507b8..1394ec6e 100644 --- a/worker/docker.go +++ b/worker/docker.go @@ -110,10 +110,12 @@ func overridePipelineImages(defaultImage string) error { // First check if the defaultImage is a valid JSON string var pipelineOverrides map[string]string - if strings.HasPrefix(defaultImage, "{") && strings.HasSuffix(defaultImage, "}") { + if strings.HasPrefix(defaultImage, "{") || strings.HasSuffix(defaultImage, "}") { // Parse the JSON string for custom pipeline images override - if err := json.Unmarshal([]byte(defaultImage), &pipelineOverrides); err != nil { - return fmt.Errorf("failed to parse pipeline overrides JSON: %w", err) + err := json.Unmarshal([]byte(defaultImage), &pipelineOverrides) + if err != nil { + slog.Error("Malformed JSON in pipeline override", "error", err, "json", defaultImage) + return err } // Override the pipelineToImage and livePipelineToImage mappings diff --git a/worker/docker_test.go b/worker/docker_test.go index 8c405383..ed15cbd2 100644 --- a/worker/docker_test.go +++ b/worker/docker_test.go @@ -813,7 +813,7 @@ func TestDockerManager_overridePipelineImages(t *testing.T) { { name: "NoOverrideFallback", inputJSON: `{"segment-anything-2": "custom-image:1.0"}`, - pipeline: "video-to-video", + pipeline: "streamdiffusion", expectedImage: "default-image", expectError: false, }, @@ -821,12 +821,12 @@ func TestDockerManager_overridePipelineImages(t *testing.T) { name: "EmptyJSON", inputJSON: `{}`, pipeline: "segment-anything-2", - expectedImage: "default-image", + expectedImage: "custom-image:1.0", expectError: false, }, { name: "MalformedJSON", - inputJSON: `{"segment-anything-2": "custom-image:1.0`, + inputJSON: `{"segment-anything-2": "custom-image:1.0"`, pipeline: "segment-anything-2", expectError: true, }, From 878d848ea18c3357505f690692b10848ccd7ef28 Mon Sep 17 00:00:00 2001 From: RUFFY-369 Date: Mon, 2 Dec 2024 14:27:12 +0530 Subject: [PATCH 5/7] test:add regular string input case --- worker/docker_test.go | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/worker/docker_test.go b/worker/docker_test.go index ed15cbd2..fc1b81c6 100644 --- a/worker/docker_test.go +++ b/worker/docker_test.go @@ -830,6 +830,13 @@ func TestDockerManager_overridePipelineImages(t *testing.T) { pipeline: "segment-anything-2", expectError: true, }, + { + name: "RegularStringInput", + inputJSON: "", + pipeline: "image-to-video", + expectedImage: "default-image", + expectError: false, + }, } for _, tt := range tests { From 50271f7997915a2b849c0b0f62d0d5f745e38c9c Mon Sep 17 00:00:00 2001 From: RUFFY-369 Date: Tue, 3 Dec 2024 14:26:13 +0530 Subject: [PATCH 6/7] refactor: update docstring placement to follow idiomatic Go convention --- worker/docker.go | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/worker/docker.go b/worker/docker.go index 1394ec6e..b57f2963 100644 --- a/worker/docker.go +++ b/worker/docker.go @@ -99,15 +99,14 @@ type DockerManager struct { mu *sync.Mutex } +// overridePipelineImages function parses a JSON string containing pipeline-to-image mappings and overrides the default mappings if valid. +// It updates the `pipelineToImage` and `livePipelineToImage` maps with custom images. +// Parameters: +// - defaultImage: A string that can either be containerImage name or a JSON string with overrides for pipeline-to-image mappings. +// +// Returns: +// - error: An error if the JSON parsing fails or if the mapping is not found in existing maps else `nil`. func overridePipelineImages(defaultImage string) error { - // overridePipelineImages function parses a JSON string containing pipeline-to-image mappings and overrides the default mappings if valid. - // It updates the `pipelineToImage` and `livePipelineToImage` maps with custom images. - // Parameters: - // - defaultImage: A string that can either be containerImage name or a JSON string with overrides for pipeline-to-image mappings. - // - // Returns: - // - error: An error if the JSON parsing fails or if the mapping is not found in existing maps else `nil`. - // First check if the defaultImage is a valid JSON string var pipelineOverrides map[string]string if strings.HasPrefix(defaultImage, "{") || strings.HasSuffix(defaultImage, "}") { From ddfb96b81cb2eae484973ddf2345175892fcd1f8 Mon Sep 17 00:00:00 2001 From: RUFFY-369 Date: Tue, 3 Dec 2024 15:35:57 +0530 Subject: [PATCH 7/7] refactor:optimise parsing func --- worker/docker.go | 43 ++++++++++++++++++++++--------------------- 1 file changed, 22 insertions(+), 21 deletions(-) diff --git a/worker/docker.go b/worker/docker.go index b57f2963..3ca710fe 100644 --- a/worker/docker.go +++ b/worker/docker.go @@ -99,6 +99,23 @@ type DockerManager struct { mu *sync.Mutex } +// updatePipelineMappings updates the specified mapping with pipeline to image overriding. +// It logs a warning if a pipeline is not found in the given mapping. +// +// Parameters: +// - overrides: A map of pipeline names to custom image names. +// - mapping: The map to be updated with the provided overrides. +// - mapName: The name of the map (used for logging purposes). +func updatePipelineMappings(overrides map[string]string, mapping map[string]string, mapName string) { + for pipeline, image := range overrides { + if _, exists := mapping[pipeline]; exists { + mapping[pipeline] = image + } else { + slog.Warn("Pipeline not found in map", "map", mapName, "pipeline", pipeline) + } + } +} + // overridePipelineImages function parses a JSON string containing pipeline-to-image mappings and overrides the default mappings if valid. // It updates the `pipelineToImage` and `livePipelineToImage` maps with custom images. // Parameters: @@ -107,30 +124,14 @@ type DockerManager struct { // Returns: // - error: An error if the JSON parsing fails or if the mapping is not found in existing maps else `nil`. func overridePipelineImages(defaultImage string) error { - // First check if the defaultImage is a valid JSON string - var pipelineOverrides map[string]string if strings.HasPrefix(defaultImage, "{") || strings.HasSuffix(defaultImage, "}") { - // Parse the JSON string for custom pipeline images override - err := json.Unmarshal([]byte(defaultImage), &pipelineOverrides) - if err != nil { - slog.Error("Malformed JSON in pipeline override", "error", err, "json", defaultImage) + var pipelineOverrides map[string]string + if err := json.Unmarshal([]byte(defaultImage), &pipelineOverrides); err != nil { + slog.Error("Error parsing JSON", "error", err) return err } - - // Override the pipelineToImage and livePipelineToImage mappings - for pipeline, image := range pipelineOverrides { - if _, isPresent := pipelineToImage[pipeline]; isPresent { - pipelineToImage[pipeline] = image - } else { - slog.Warn("Pipeline not found in pipelineToImage map", "pipeline", pipeline) - } - - if _, isPresent := livePipelineToImage[pipeline]; isPresent { - livePipelineToImage[pipeline] = image - } else { - slog.Warn("Pipeline not found in livePipelineToImage map", "pipeline", pipeline) - } - } + updatePipelineMappings(pipelineOverrides, pipelineToImage, "pipelineToImage") + updatePipelineMappings(pipelineOverrides, livePipelineToImage, "livePipelineToImage") } return nil }