Skip to content

Commit

Permalink
Add getters for serviecs that have defined build, extends and depends_on
Browse files Browse the repository at this point in the history
Signed-off-by: jhrotko <joana.hrotko@docker.com>

Add capanilities,gpu,tpu count per service

Signed-off-by: jhrotko <joana.hrotko@docker.com>

Apply Filter to Services

Signed-off-by: jhrotko <joana.hrotko@docker.com>

cleanup

Signed-off-by: jhrotko <joana.hrotko@docker.com>
  • Loading branch information
jhrotko committed Feb 7, 2024
1 parent 3c5150f commit 64db8a6
Show file tree
Hide file tree
Showing 4 changed files with 201 additions and 0 deletions.
53 changes: 53 additions & 0 deletions types/project.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ import (
"github.com/distribution/reference"
"github.com/mitchellh/copystructure"
godigest "github.com/opencontainers/go-digest"
"golang.org/x/exp/maps"
"golang.org/x/sync/errgroup"
"gopkg.in/yaml.v3"
)
Expand Down Expand Up @@ -119,6 +120,58 @@ func (p *Project) ConfigNames() []string {
return names
}

func (p *Project) ServicesBuild() []string {
servicesBuild := p.Services.Filter(func(s ServiceConfig) bool {
return s.Build != nil && s.Build.Context != ""
})
return maps.Keys(servicesBuild)
}

func (p *Project) ServicesExtends() []string {
servicesExtends := p.Services.Filter(func(s ServiceConfig) bool {
return s.Extends != nil && *s.Extends != (ExtendsConfig{})
})
return maps.Keys(servicesExtends)
}

func (p *Project) ServicesDependsOn() []string {
servicesDependsOn := p.Services.Filter(func(s ServiceConfig) bool {
return len(s.DependsOn) > 0
})
return maps.Keys(servicesDependsOn)
}

func (p *Project) ServicesCapabilities() ([]string, []string, []string) {
capabilities := []string{}
gpu := []string{}
tpu := []string{}
for _, service := range p.Services {
deploy := service.Deploy
if deploy == nil {
continue
}
reservation := deploy.Resources.Reservations
if reservation == nil {
continue
}
devices := reservation.Devices
for _, d := range devices {
if len(d.Capabilities) > 0 {
capabilities = append(capabilities, service.Name)
}
for _, c := range d.Capabilities {
if c == "gpu" {
gpu = append(gpu, service.Name)
} else if c == "tpu" {
tpu = append(tpu, service.Name)
}
}
}
}

return utils.RemoveDuplicates(capabilities), utils.RemoveDuplicates(gpu), utils.RemoveDuplicates(tpu)
}

// GetServices retrieve services by names, or return all services if no name specified
func (p *Project) GetServices(names ...string) (Services, error) {
if len(names) == 0 {
Expand Down
123 changes: 123 additions & 0 deletions types/project_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -240,3 +240,126 @@ func TestWithServices(t *testing.T) {
assert.NilError(t, err)
assert.DeepEqual(t, seen, []string{"service_1", "service_2", "service_4"})
}

func TestBuildServices(t *testing.T) {
p := makeProject()
assert.DeepEqual(t, []string{}, p.ServicesBuild())

service, err := p.GetService("service_1")
assert.NilError(t, err)
service.Build = &BuildConfig{}
p.Services["service_1"] = service
assert.DeepEqual(t, []string{}, p.ServicesBuild())

service.Build = &BuildConfig{
Context: ".",
}
p.Services["service_1"] = service
service, err = p.GetService("service_2")
assert.NilError(t, err)
service.Build = &BuildConfig{
Context: ".",
}
p.Services["service_2"] = service
assert.DeepEqual(t, []string{"service_1", "service_2"}, p.ServicesBuild())
}

func TestExtendsNumber(t *testing.T) {
p := makeProject()
assert.DeepEqual(t, []string{}, p.ServicesExtends())

service, err := p.GetService("service_1")
assert.NilError(t, err)
service.Extends = &ExtendsConfig{}
p.Services["service_1"] = service
assert.DeepEqual(t, []string{}, p.ServicesExtends())

service.Extends = &ExtendsConfig{
File: ".",
Service: "service_2",
}
p.Services["service_1"] = service
assert.DeepEqual(t, []string{"service_1"}, p.ServicesExtends())
}

func TestServicesDependsOn(t *testing.T) {
p := makeProject()
assert.Equal(t, 3, len(p.ServicesDependsOn()))
assert.DeepEqual(t, []string{"service_2", "service_3", "service_4"}, p.ServicesDependsOn())
}

func TestServicesCapabilities(t *testing.T) {
p := makeProject()

service, err := p.GetService("service_1")
assert.NilError(t, err)
service.Deploy = &DeployConfig{}
p.Services["service_1"] = service
capabilities, gpu, tpu := p.ServicesCapabilities()
assert.DeepEqual(t, []string{}, capabilities)
assert.DeepEqual(t, []string{}, gpu)
assert.DeepEqual(t, []string{}, tpu)

service.Deploy = &DeployConfig{
Resources: Resources{
Reservations: &Resource{},
},
}
p.Services["service_1"] = service
capabilities, gpu, tpu = p.ServicesCapabilities()
assert.DeepEqual(t, []string{}, capabilities)
assert.DeepEqual(t, []string{}, gpu)
assert.DeepEqual(t, []string{}, tpu)

service.Deploy = &DeployConfig{
Resources: Resources{
Reservations: &Resource{
Devices: []DeviceRequest(nil),
},
},
}
p.Services["service_1"] = service
capabilities, gpu, tpu = p.ServicesCapabilities()
assert.DeepEqual(t, []string{}, capabilities)
assert.DeepEqual(t, []string{}, gpu)
assert.DeepEqual(t, []string{}, tpu)

service.Deploy = &DeployConfig{
Resources: Resources{
Reservations: &Resource{
Devices: []DeviceRequest{
{
Capabilities: []string{"gpu", "tpu"},
},
},
},
},
}
p.Services["service_1"] = service
capabilities, gpu, tpu = p.ServicesCapabilities()
assert.DeepEqual(t, []string{"service_1"}, capabilities)
assert.DeepEqual(t, []string{"service_1"}, gpu)
assert.DeepEqual(t, []string{"service_1"}, tpu)

service, err = p.GetService("service_2")
assert.NilError(t, err)
service.Deploy = &DeployConfig{
Resources: Resources{
Reservations: &Resource{
Devices: []DeviceRequest{
{
Capabilities: []string{"tpu"},
},
{
Capabilities: []string{"tpu"},
},
},
},
},
}
p.Services["service_2"] = service
capabilities, gpu, tpu = p.ServicesCapabilities()
assert.DeepEqual(t, []string{"service_1", "service_2"}, capabilities)
assert.DeepEqual(t, []string{"service_1"}, gpu)
assert.DeepEqual(t, []string{"service_1", "service_2"}, tpu)
}
10 changes: 10 additions & 0 deletions types/services.go
Original file line number Diff line number Diff line change
Expand Up @@ -33,3 +33,13 @@ func (s Services) GetProfiles() []string {
}
return profiles
}

func (s Services) Filter(predicate func(ServiceConfig) bool) Services {
services := Services{}
for name, service := range s {
if predicate(service) {
services[name] = service
}
}
return services
}
15 changes: 15 additions & 0 deletions utils/collectionutils.go
Original file line number Diff line number Diff line change
Expand Up @@ -51,3 +51,18 @@ func ArrayContains[T comparable](source []T, toCheck []T) bool {
}
return true
}

func RemoveDuplicates[T comparable](slice []T) []T {
// Create a map to store unique elements
seen := make(map[T]bool)
result := []T{}

// Loop through the slice, adding elements to the map if they haven't been seen before
for _, val := range slice {
if _, ok := seen[val]; !ok {
seen[val] = true
result = append(result, val)
}
}
return result
}

0 comments on commit 64db8a6

Please sign in to comment.