Check resources using a command

README.md

@@ -66,7 +66,11 @@ Below is an example `config.jsonc`:
   "MaxTimeToWaitForServiceToCloseConnectionBeforeGivingUpSeconds": 1200,
   "ShutDownAfterInactivitySeconds": 120,
   "ResourcesAvailable": {
-    "VRAM-GPU-1": 24000,
+    "VRAM-GPU-1": {
+        "Amount": 24000,
+        "CheckCommand": "nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i 0",
+        "CheckIntervalMilliseconds": 1000,
+    }
     "RAM": 32000,
   },
   "Services": [

config.go

@@ -82,9 +82,9 @@ type Config struct {
 	ShutDownAfterInactivitySeconds                                uint
 	MaxTimeToWaitForServiceToCloseConnectionBeforeGivingUpSeconds *uint
 	OutputServiceLogs                                             *bool
-	DefaultServiceUrl                                             *string         `json:"DefaultServiceUrl"`
-	Services                                                      []ServiceConfig `json:"Services"`
-	ResourcesAvailable                                            map[string]int  `json:"ResourcesAvailable"`
+	DefaultServiceUrl                                             *string                      `json:"DefaultServiceUrl"`
+	Services                                                      []ServiceConfig              `json:"Services"`
+	ResourcesAvailable                                            map[string]ResourceAvailable `json:"ResourcesAvailable"`
 	OpenAiApi                                                     OpenAiApi
 	ManagementApi                                                 ManagementApi
 }
@@ -110,13 +110,57 @@ type ServiceConfig struct {
 	ServiceUrl                      *ServiceUrlOption `json:"ServiceUrl,omitempty"`
 	ResourceRequirements            map[string]int    `json:"ResourceRequirements"`
 }
+type ResourceAvailable struct {
+	Amount                    int
+	CheckCommand              string
+	CheckIntervalMilliseconds uint
+}
+
+// Accepts either a JSON number or an object with {Amount, CheckCommand}.
+func (r *ResourceAvailable) UnmarshalJSON(data []byte) error {
+	trimmed := bytes.TrimSpace(data)
+	if bytes.Equal(trimmed, []byte("null")) {
+		*r = ResourceAvailable{}
+		return nil
+	}
+
+	var asInt int
+	err := json.Unmarshal(trimmed, &asInt)
+	if err == nil {
+		*r = ResourceAvailable{Amount: asInt}
+		return nil
+	}
+
+	var dto struct {
+		Amount                    int    `json:"Amount"`
+		CheckCommand              string `json:"CheckCommand"`
+		CheckIntervalMilliseconds uint   `json:"CheckIntervalMilliseconds"`
+	}
+
+	dec := json.NewDecoder(bytes.NewReader(trimmed))
+	dec.DisallowUnknownFields()
+	err = dec.Decode(&dto)
+
+	if err == nil && !(dto.Amount == 0 && dto.CheckCommand == "") {
+		if dto.CheckIntervalMilliseconds == 0 {
+			dto.CheckIntervalMilliseconds = 1000
+		}
+		*r = ResourceAvailable{Amount: dto.Amount, CheckCommand: dto.CheckCommand, CheckIntervalMilliseconds: dto.CheckIntervalMilliseconds}
+		return nil
+	}
+
+	if err == nil {
+		err = errors.New("missing both Amount and CheckCommand fields")
+	}
+	return fmt.Errorf("each entry in ResourcesAvailable must be an integer or an object with at least one of the fields: \"Amount\", \"CheckCommand\", e.g. \"ResourceAvailable: {\"RAM\": {\"Amount\": 1, \"CheckCommand\": \"echo 1\"}} %v", err)
+}
 
-// UnmarshalJSON implements custom unmarshaling for ServiceConfig to handle ServiceUrl properly
+// UnmarshalJSON implements custom unmarshaling for ServiceConfig to handle ServiceUrl and ResourcesAvailable properly
 func (sc *ServiceConfig) UnmarshalJSON(data []byte) error {
-	// Create a type alias to avoid infinite recursion
 	type Alias ServiceConfig
 	aux := &struct {
-		ServiceUrl json.RawMessage `json:"ServiceUrl"`
+		ServiceUrl         json.RawMessage              `json:"ServiceUrl"`
+		ResourcesAvailable map[string]ResourceAvailable `json:"ResourcesAvailable"`
 		*Alias
 	}{
 		Alias: (*Alias)(sc),

config_test.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"github.com/stretchr/testify/assert"
 	"strings"
 	"testing"
 )
@@ -25,10 +26,11 @@ func loadConfigFromString(t *testing.T, jsonStr string) (Config, error) {
 
 func TestValidConfigMinimal(t *testing.T) {
 	t.Parallel()
-	_, err := loadConfigFromString(t, `{
+	config, err := loadConfigFromString(t, `{
 		"ResourcesAvailable": {
 			"RAM": 10000, //Random access memory!
 			"VRAM-GPU-1": 20000, /* Wow, VRAM! */
+            "VRAM-GPU-2": {"Amount": 1000, "CheckCommand": "Look, I am an object!", "CheckIntervalMilliseconds": 4000},
 		},
         /* A multi-line comment
            to make sure JSONc works
@@ -45,13 +47,33 @@ func TestValidConfigMinimal(t *testing.T) {
 			{
 				"Name": "serviceB",
 				"ListenPort": "8081",
-				"Command": "/bin/echo"
+				"Command": "/echo/bin"
 			}
 		]
 	}`)
 	if err != nil {
 		t.Fatalf("did not expect an error but got: %v", err)
 	}
+	assert.Contains(t, config.ResourcesAvailable, "RAM")
+	assert.Contains(t, config.ResourcesAvailable, "VRAM-GPU-1")
+	assert.Contains(t, config.ResourcesAvailable, "VRAM-GPU-2")
+	assert.Len(t, config.ResourcesAvailable, 3)
+	assert.Equal(t, 10000, config.ResourcesAvailable["RAM"].Amount)
+	assert.Equal(t, "", config.ResourcesAvailable["RAM"].CheckCommand)
+	assert.Equal(t, 20000, config.ResourcesAvailable["VRAM-GPU-1"].Amount)
+	assert.Equal(t, "", config.ResourcesAvailable["VRAM-GPU-1"].CheckCommand)
+	assert.Equal(t, 1000, config.ResourcesAvailable["VRAM-GPU-2"].Amount)
+	assert.Equal(t, "Look, I am an object!", config.ResourcesAvailable["VRAM-GPU-2"].CheckCommand)
+	assert.Equal(t, uint(4000), config.ResourcesAvailable["VRAM-GPU-2"].CheckIntervalMilliseconds)
+	assert.Equal(t, "7070", config.OpenAiApi.ListenPort)
+	assert.Equal(t, "", config.ManagementApi.ListenPort)
+	assert.Len(t, config.Services, 2)
+	assert.Equal(t, "serviceA", config.Services[0].Name)
+	assert.Equal(t, "8080", config.Services[0].ListenPort)
+	assert.Equal(t, "/bin/echo", config.Services[0].Command)
+	assert.Equal(t, "serviceB", config.Services[1].Name)
+	assert.Equal(t, "8081", config.Services[1].ListenPort)
+	assert.Equal(t, "/echo/bin", config.Services[1].Command)
 }
 
 func TestDuplicateServiceNames(t *testing.T) {
@@ -900,3 +922,54 @@ func TestJsonWithInvalidDataAfterObject(t *testing.T) {
 		"extra data after the first JSON object",
 	})
 }
+
+func TestJsonWithExtraDataInResourcesAvailable(t *testing.T) {
+	t.Parallel()
+	_, err := loadConfigFromString(t, `{
+		"ResourcesAvailable": { "RAM": {"Amount": 10, "Foo": "Bar"} },
+}`)
+	checkExpectedErrorMessages(t, err, []string{
+		"each entry in ResourcesAvailable must be an integer or an object with at least one of the fields",
+	})
+}
+
+func TestJsonWithNonIntAmount(t *testing.T) {
+	t.Parallel()
+	_, err := loadConfigFromString(t, `{
+		"ResourcesAvailable": { "RAM": {"Amount": "Bar",} },
+	}`)
+	checkExpectedErrorMessages(t, err, []string{
+		"each entry in ResourcesAvailable must be an integer or an object with at least one of the fields",
+	})
+}
+
+func TestJsonWithNoAmountOrCommand(t *testing.T) {
+	t.Parallel()
+	_, err := loadConfigFromString(t, `{
+		"ResourcesAvailable": { "RAM": {} },
+	}`)
+	checkExpectedErrorMessages(t, err, []string{
+		"each entry in ResourcesAvailable must be an integer or an object with at least one of the fields",
+	})
+}
+
+func TestJsonWithCheckCommand(t *testing.T) {
+	t.Parallel()
+	config, err := loadConfigFromString(t, `{
+		"ResourcesAvailable": { "RAM": {"CheckCommand": "Bar",} },
+	}`)
+	if err != nil {
+		t.Fatalf("did not expect an error but got: %v", err)
+	}
+	assert.Equal(t, "Bar", config.ResourcesAvailable["RAM"].CheckCommand)
+	assert.Equal(t, uint(1000), config.ResourcesAvailable["RAM"].CheckIntervalMilliseconds)
+}
+func TestJsonWithNonIntegerInCheckInterval(t *testing.T) {
+	t.Parallel()
+	_, err := loadConfigFromString(t, `{
+		"ResourcesAvailable": { "CheckCommand": "Bar", "CheckFrequencyMilliseconds": "string" },
+	}`)
+	checkExpectedErrorMessages(t, err, []string{
+		"each entry in ResourcesAvailable must be an integer or an object with at least one of the fields",
+	})
+}

go.mod

@@ -9,6 +9,6 @@ require github.com/tidwall/jsonc v0.3.2
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/stretchr/testify v1.10.0 // indirect
+	github.com/stretchr/testify v1.11.1 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -6,6 +6,8 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
 github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
+github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
 github.com/tidwall/jsonc v0.3.2 h1:ZTKrmejRlAJYdn0kcaFqRAKlxxFIC21pYq8vLa4p2Wc=
 github.com/tidwall/jsonc v0.3.2/go.mod h1:dw+3CIxqHi+t8eFSpzzMlcVYxKp08UP5CD8/uSFCyJE=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

main.go

@@ -36,9 +36,10 @@ type RunningService struct {
 }
 
 type ResourceManager struct {
-	serviceMutex    *sync.Mutex
-	resourcesInUse  map[string]int
-	runningServices map[string]RunningService
+	serviceMutex       *sync.Mutex
+	resourcesInUse     map[string]int
+	resourcesAvailable map[string]*int
+	runningServices    map[string]RunningService
 }
 type OpenAiApiModels struct {
 	Object string           `json:"object"`
@@ -149,11 +150,25 @@ func main() {
 	}
 
 	resourceManager = ResourceManager{
-		resourcesInUse:  make(map[string]int),
-		runningServices: make(map[string]RunningService),
-		serviceMutex:    &sync.Mutex{},
+		resourcesInUse:     make(map[string]int),
+		resourcesAvailable: make(map[string]*int),
+		runningServices:    make(map[string]RunningService),
+		serviceMutex:       &sync.Mutex{},
+	}
+
+	for name, resource := range config.ResourcesAvailable {
+		//Using int reference to avoid having a lock for reading from the map
+		resourceManager.resourcesAvailable[name] = new(int)
+		*resourceManager.resourcesAvailable[name] = resource.Amount
+		if resource.CheckCommand != "" {
+			go monitorResourceAvailability(
+				name,
+				resource.CheckCommand,
+				time.Duration(resource.CheckIntervalMilliseconds)*time.Millisecond,
+				&resourceManager,
+			)
+		}
 	}
-
 	for _, service := range config.Services {
 		if service.ListenPort != "" {
 			go startProxy(service)
@@ -907,13 +922,19 @@ func findEarliestLastUsedServiceUsingResource(requestingService string, missingR
 
 func findFirstMissingResourceWhenServiceMutexIsLocked(resourceRequirements map[string]int, requestingService string, outputError bool) *string {
 	for resource, amount := range resourceRequirements {
-		if resourceManager.resourcesInUse[resource]+amount > config.ResourcesAvailable[resource] {
+		var enoughOfResource bool
+		if config.ResourcesAvailable[resource].CheckCommand == "" {
+			enoughOfResource = resourceManager.resourcesInUse[resource]+amount > *resourceManager.resourcesAvailable[resource]
+		} else { //Ignore resources in use if available amount is dynamically calculated
+			enoughOfResource = amount > *resourceManager.resourcesAvailable[resource]
+		}
+		if !enoughOfResource {
 			if outputError {
 				log.Printf(
 					"[%s] Not enough %s to start. Total: %d, In use: %d, Required: %d",
 					requestingService,
 					resource,
-					config.ResourcesAvailable[resource],
+					resourceManager.resourcesAvailable[resource],
 					resourceManager.resourcesInUse[resource],
 					amount,
 				)