You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airavata.apache.org by di...@apache.org on 2026/05/20 05:06:38 UTC
(airavata-custos) branch master updated (bf4d2e148 -> 3d102159e)
This is an automated email from the ASF dual-hosted git repository.
DImuthuUpe pushed a change to branch master
in repository https://gitbox.apache.org/repos/asf/airavata-custos.git
from bf4d2e148 Revert last 3 commits
new 8cf058ff1 Initial framework in integrate connectors to core with pub sub
new 983beccd1 Creating a slurm account when an compute allocation was created
new 9b7c6e544 Fixing minor misconfiguration
new 6b7e6daa6 Adding SLURM association when an allocation membership is created
new 3b0bf6845 Moving member-specific quota to a seperate data model and implementing slurm subscriber
new 13ef537fd Updating account resource limits when the allocation resource mapping is created
new 3d102159e Enabling member level allocation restriction
The 7 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
cmd/server/main.go | 7 +-
connectors/SLURM/Association-Mapper/Makefile | 18 --
connectors/SLURM/Association-Mapper/README.md | 56 ++--
connectors/SLURM/Association-Mapper/go.mod | 3 -
.../internal/operations/accounts.go | 8 +-
.../internal/operations/accounts_test.go | 8 +-
.../internal/operations/associations.go | 6 +-
.../internal/operations/associations_test.go | 7 +-
.../internal/operations/client.go | 20 +-
.../internal/operations/client_test.go | 4 +-
.../internal/subscribers/account.go | 121 ++++++++
.../internal/subscribers/members.go | 154 ++++++++++
.../internal/subscribers/subscriber.go | 28 ++
connectors/SLURM/Association-Mapper/main.go | 25 --
.../SLURM/Association-Mapper/pkg/smapper/loader.go | 27 ++
docs/API-Docs.md | 341 +++++++++++++++++++--
internal/connectors/loader.go | 21 ++
...n.sql => 000010_compute_cluster_users.down.sql} | 2 +-
....up.sql => 000010_compute_cluster_users.up.sql} | 24 +-
...ocation_membership_resource_overrides.down.sql} | 5 +-
...allocation_membership_resource_overrides.up.sql | 36 +++
...e_allocation_resource_mapping_amounts.down.sql} | 4 +-
...ute_allocation_resource_mapping_amounts.up.sql} | 6 +-
..._membership_resource_override_amounts.down.sql} | 4 +-
...on_membership_resource_override_amounts.up.sql} | 6 +-
internal/server/server.go | 192 ++++++++++--
...llocation_membership_resource_override_store.go | 124 ++++++++
.../store/compute_allocation_membership_store.go | 11 +-
.../compute_allocation_resource_mapping_store.go | 33 +-
internal/store/compute_cluster_user_store.go | 116 +++++++
internal/store/store.go | 43 +++
...ation_membership_resource_override_subscribe.go | 67 ++++
pkg/events/compute_cluster_user_subscribe.go | 64 ++++
pkg/events/types.go | 14 +
pkg/models/allocation.go | 26 +-
pkg/service/compute_allocation_membership.go | 27 --
...pute_allocation_membership_resource_override.go | 197 ++++++++++++
pkg/service/compute_allocation_resource_mapping.go | 50 ++-
pkg/service/compute_cluster.go | 6 +-
pkg/service/compute_cluster_user.go | 179 +++++++++++
pkg/service/service.go | 8 +
41 files changed, 1886 insertions(+), 212 deletions(-)
delete mode 100644 connectors/SLURM/Association-Mapper/Makefile
delete mode 100644 connectors/SLURM/Association-Mapper/go.mod
create mode 100644 connectors/SLURM/Association-Mapper/internal/subscribers/account.go
create mode 100644 connectors/SLURM/Association-Mapper/internal/subscribers/members.go
create mode 100644 connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go
delete mode 100644 connectors/SLURM/Association-Mapper/main.go
create mode 100644 connectors/SLURM/Association-Mapper/pkg/smapper/loader.go
create mode 100644 internal/connectors/loader.go
copy internal/db/migrations/{000002_compute_clusters.down.sql => 000010_compute_cluster_users.down.sql} (94%)
copy internal/db/migrations/{000006_compute_allocation_diffs.up.sql => 000010_compute_cluster_users.up.sql} (53%)
copy internal/db/migrations/{000004_compute_allocation_resource_mappings.down.sql => 000011_compute_allocation_membership_resource_overrides.down.sql} (80%)
create mode 100644 internal/db/migrations/000011_compute_allocation_membership_resource_overrides.up.sql
copy internal/db/migrations/{000004_compute_allocation_resource_mappings.down.sql => 000012_compute_allocation_resource_mapping_amounts.down.sql} (87%)
copy internal/db/migrations/{000001_initial_schema.down.sql => 000012_compute_allocation_resource_mapping_amounts.up.sql} (78%)
copy internal/db/migrations/{000004_compute_allocation_resource_mappings.down.sql => 000013_compute_allocation_membership_resource_override_amounts.down.sql} (80%)
copy internal/db/migrations/{000001_initial_schema.down.sql => 000013_compute_allocation_membership_resource_override_amounts.up.sql} (75%)
create mode 100644 internal/store/compute_allocation_membership_resource_override_store.go
create mode 100644 internal/store/compute_cluster_user_store.go
create mode 100644 pkg/events/compute_allocation_membership_resource_override_subscribe.go
create mode 100644 pkg/events/compute_cluster_user_subscribe.go
create mode 100644 pkg/service/compute_allocation_membership_resource_override.go
create mode 100644 pkg/service/compute_cluster_user.go
(airavata-custos) 01/07: Initial framework in integrate connectors to core with pub sub
Posted by di...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
DImuthuUpe pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/airavata-custos.git
commit 8cf058ff173623d7707c6a7654b16f4e7eaf2f67
Author: DImuthuUpe <di...@gmail.com>
AuthorDate: Mon May 18 00:38:34 2026 -0400
Initial framework in integrate connectors to core with pub sub
---
cmd/server/main.go | 5 +++
connectors/SLURM/Association-Mapper/Makefile | 18 ----------
connectors/SLURM/Association-Mapper/README.md | 42 +++++-----------------
connectors/SLURM/Association-Mapper/go.mod | 3 --
.../internal/subscribers/account.go | 16 +++++++++
.../internal/subscribers/subscriber.go | 22 ++++++++++++
connectors/SLURM/Association-Mapper/main.go | 25 -------------
.../SLURM/Association-Mapper/pkg/smapper/loader.go | 11 ++++++
internal/connectors/loader.go | 20 +++++++++++
9 files changed, 82 insertions(+), 80 deletions(-)
diff --git a/cmd/server/main.go b/cmd/server/main.go
index 978d5cc92..ea0ebd697 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -29,6 +29,7 @@ import (
"syscall"
"time"
+ "github.com/apache/airavata-custos/internal/connectors"
"github.com/apache/airavata-custos/internal/db"
"github.com/apache/airavata-custos/internal/server"
"github.com/apache/airavata-custos/pkg/events"
@@ -72,6 +73,10 @@ func run() error {
// Create a new event bus instance to async messaging between service and connectors
eventBus := events.New()
+ if err := connectors.LoadConnectors(eventBus); err != nil {
+ return err
+ }
+
svc := service.New(database, eventBus)
handler := server.LoggingMiddleware(server.New(svc))
diff --git a/connectors/SLURM/Association-Mapper/Makefile b/connectors/SLURM/Association-Mapper/Makefile
deleted file mode 100644
index 32efab927..000000000
--- a/connectors/SLURM/Association-Mapper/Makefile
+++ /dev/null
@@ -1,18 +0,0 @@
-.PHONY: build test run tidy clean
-
-BIN := bin/association-mapper
-
-build:
- go build -o $(BIN) .
-
-run: build
- ./$(BIN)
-
-test:
- go test ./...
-
-tidy:
- go mod tidy
-
-clean:
- rm -rf bin
diff --git a/connectors/SLURM/Association-Mapper/README.md b/connectors/SLURM/Association-Mapper/README.md
index 667af800f..9d9bfc6f0 100644
--- a/connectors/SLURM/Association-Mapper/README.md
+++ b/connectors/SLURM/Association-Mapper/README.md
@@ -2,6 +2,8 @@
SLURM association creation logic lives in this plugin. It is triggered when the allocation manager has processed an allocation request and released it to downstream handlers. It talks to `slurmrestd` to manage accounts, associations, and TRES limits.
+This package is part of the root `github.com/apache/airavata-custos` module.
+
## Prerequisites
- Go **1.24+**
@@ -11,46 +13,18 @@ SLURM association creation logic lives in this plugin. It is triggered when the
```
.
-├── main.go # entry point
-├── internal/operations/ # slurmrestd client + accounts/associations/TRES
-├── go.mod
-└── Makefile
-```
-
-Module path: `github.com/apache/airavata-custos/connectors/SLURM/Association-Mapper`.
-
-## Build
-
-```bash
-# from this directory
-make build # produces bin/association-mapper
-# or directly:
-go build -o bin/association-mapper .
+├── internal/operations/ # slurmrestd client + accounts/associations/TRES
+└── pkg/operations/
```
-## Run
-
-```bash
-make run # build, then ./bin/association-mapper
-```
-
-The service starts, logs `association-mapper started`, and blocks until it receives `SIGINT` or `SIGTERM`.
+Import path: `github.com/apache/airavata-custos/connectors/SLURM/Association-Mapper/internal/operations`.
## Test
```bash
-make test # go test ./...
-go vet ./... # static checks
+# from the repository root
+go test ./connectors/SLURM/Association-Mapper/...
+go vet ./connectors/SLURM/Association-Mapper/...
```
Tests are hermetic and use `httptest` — no live `slurmrestd` required.
-
-## Common make targets
-
-| Target | Description |
-|---------|--------------------------------------|
-| `build` | Compile the binary into `bin/` |
-| `run` | Build and run |
-| `test` | Run all unit tests |
-| `tidy` | `go mod tidy` |
-| `clean` | Remove `bin/` |
diff --git a/connectors/SLURM/Association-Mapper/go.mod b/connectors/SLURM/Association-Mapper/go.mod
deleted file mode 100644
index 0e1832e7c..000000000
--- a/connectors/SLURM/Association-Mapper/go.mod
+++ /dev/null
@@ -1,3 +0,0 @@
-module github.com/apache/airavata-custos/connectors/SLURM/Association-Mapper
-
-go 1.24.0
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/account.go b/connectors/SLURM/Association-Mapper/internal/subscribers/account.go
new file mode 100644
index 000000000..9c2f29c6f
--- /dev/null
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/account.go
@@ -0,0 +1,16 @@
+package subscribers
+
+import "github.com/apache/airavata-custos/pkg/models"
+import "log/slog"
+
+func (a *AssociationSubscriber) SubscribeToComputeAccountCreation(computeAccount models.ComputeAllocation) {
+ slog.Info("Received compute account creation event", "account", computeAccount)
+}
+
+func (a *AssociationSubscriber) SubscribeToComputeAccountDeletion(computeAccount models.ComputeAllocation) {
+ slog.Info("Received compute account deletion event", "account", computeAccount)
+}
+
+func (a *AssociationSubscriber) SubscribeToComputeAccountUpdate(computeAccount models.ComputeAllocation) {
+ slog.Info("Received compute account update event", "account", computeAccount)
+}
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go b/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go
new file mode 100644
index 000000000..129bfcff9
--- /dev/null
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go
@@ -0,0 +1,22 @@
+package subscribers
+
+import client "github.com/apache/airavata-custos/connectors/SLURM/Association-Mapper/internal/operations"
+import "github.com/apache/airavata-custos/pkg/events"
+
+type AssociationSubscriber struct {
+ slurmClient *client.Client
+ eventBus *events.Bus
+}
+
+func NewAssociationSubscriber(slurmClient *client.Client, eventBus *events.Bus) *AssociationSubscriber {
+ return &AssociationSubscriber{
+ slurmClient: slurmClient,
+ eventBus: eventBus,
+ }
+}
+
+func (a *AssociationSubscriber) RegisterSubscribers() {
+ a.eventBus.SubscribeComputeAllocationCreated(a.SubscribeToComputeAccountCreation)
+ a.eventBus.SubscribeComputeAllocationDeleted(a.SubscribeToComputeAccountDeletion)
+ a.eventBus.SubscribeComputeAllocationUpdated(a.SubscribeToComputeAccountUpdate)
+}
diff --git a/connectors/SLURM/Association-Mapper/main.go b/connectors/SLURM/Association-Mapper/main.go
deleted file mode 100644
index d04dcad1e..000000000
--- a/connectors/SLURM/Association-Mapper/main.go
+++ /dev/null
@@ -1,25 +0,0 @@
-// Package main is the entry point for the SLURM Association-Mapper connector.
-//
-// It consumes allocation events released by the allocation manager and
-// materializes them as SLURM associations via slurmrestd.
-package main
-
-import (
- "context"
- "log/slog"
- "os"
- "os/signal"
- "syscall"
-)
-
-func main() {
- logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo}))
- slog.SetDefault(logger)
-
- ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
- defer stop()
-
- logger.Info("association-mapper started")
- <-ctx.Done()
- logger.Info("association-mapper stopped")
-}
diff --git a/connectors/SLURM/Association-Mapper/pkg/smapper/loader.go b/connectors/SLURM/Association-Mapper/pkg/smapper/loader.go
new file mode 100644
index 000000000..9c172b4f8
--- /dev/null
+++ b/connectors/SLURM/Association-Mapper/pkg/smapper/loader.go
@@ -0,0 +1,11 @@
+package smapper
+
+import "github.com/apache/airavata-custos/pkg/events"
+import "github.com/apache/airavata-custos/connectors/SLURM/Association-Mapper/internal/subscribers"
+import client "github.com/apache/airavata-custos/connectors/SLURM/Association-Mapper/internal/operations"
+
+func LoadConnector(eventBus *events.Bus) error {
+ slurmClient := client.New("localhost:8080", "", "") // Replace with actual SLURM client initialization.
+ subscribers.NewAssociationSubscriber(slurmClient, eventBus).RegisterSubscribers()
+ return nil
+}
diff --git a/internal/connectors/loader.go b/internal/connectors/loader.go
new file mode 100644
index 000000000..1b603abd1
--- /dev/null
+++ b/internal/connectors/loader.go
@@ -0,0 +1,20 @@
+package connectors
+
+import "github.com/apache/airavata-custos/connectors/SLURM/Association-Mapper/pkg/smapper"
+import "github.com/apache/airavata-custos/pkg/events"
+import "log/slog"
+
+func LoadConnectors(eventBus *events.Bus) error {
+
+ slog.Info("loading connectors")
+
+ slog.Info("loading SLURM Association Mapper connector")
+ err := smapper.LoadConnector(eventBus)
+ if err != nil {
+ slog.Error("failed to load SLURM Association Mapper connector", "error", err)
+ return err
+ }
+
+ slog.Info("finished loading connectors")
+ return err
+}
(airavata-custos) 07/07: Enabling member level allocation restriction
Posted by di...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
DImuthuUpe pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/airavata-custos.git
commit 3d102159ee7bd58767ad639114574e37c7440fb6
Author: DImuthuUpe <di...@gmail.com>
AuthorDate: Wed May 20 00:58:31 2026 -0400
Enabling member level allocation restriction
---
.../internal/subscribers/members.go | 8 +++----
docs/API-Docs.md | 26 +++++++++++++---------
...n_membership_resource_override_amounts.down.sql | 20 +++++++++++++++++
...ion_membership_resource_override_amounts.up.sql | 20 +++++++++++++++++
...llocation_membership_resource_override_store.go | 13 ++++++-----
pkg/models/allocation.go | 3 ++-
...pute_allocation_membership_resource_override.go | 14 ++++++++----
7 files changed, 78 insertions(+), 26 deletions(-)
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/members.go b/connectors/SLURM/Association-Mapper/internal/subscribers/members.go
index dfcc4ee70..2a087ab60 100644
--- a/connectors/SLURM/Association-Mapper/internal/subscribers/members.go
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/members.go
@@ -116,19 +116,19 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationMembershipResourceOv
grpTres := []client.TRES{}
- if allocationResource.ResourceType == "GrpTRES" {
+ if override.OverrideResourceAmount > 0 {
grpTres = append(grpTres, client.TRES{
Type: allocationResource.ResourceType,
- Count: override.OverriddenResourceAmount, // override.OverriddenResourceAmount is the SU amount, but SLURM needs the actual resource amount (e.g., number of CPU hours), so we need to convert it using the rate for the resource
+ Count: override.OverrideResourceAmount, // override.OverrideResourceAmount is the resource amount granted to the user for this resource (e.g., number of CPUs, GPUs).
})
}
grpTresMins := []client.TRES{}
- if allocationResource.ResourceType == "GrpTRESMins" {
+ if override.OverrideResourceTime > 0 {
grpTresMins = append(grpTresMins, client.TRES{
Type: allocationResource.ResourceType,
- Count: override.OverriddenResourceAmount,
+ Count: override.OverrideResourceTime, // override.OverrideResourceTime is the wall-clock time in minutes that the resource amount is granted for.
})
}
diff --git a/docs/API-Docs.md b/docs/API-Docs.md
index 325a219bc..3d9ded43c 100644
--- a/docs/API-Docs.md
+++ b/docs/API-Docs.md
@@ -1234,9 +1234,11 @@ List every per-resource override recorded against the given membership.
## Compute Allocation Membership Resource Overrides
-A `ComputeAllocationMembershipResourceOverride` records the SU amount of a
-specific resource (`ComputeAllocationResource`) that has been granted to a
-specific membership. There can be at most one override per
+A `ComputeAllocationMembershipResourceOverride` records the resource amount
+and wall-clock time of a specific resource (`ComputeAllocationResource`)
+that has been granted to a specific membership, overriding whatever the
+parent `ComputeAllocationResourceMapping` would otherwise provide. There can
+be at most one override per
`(compute_allocation_membership_id, compute_allocation_resource_id)` pair
(enforced by a unique key). Overrides are cascade-deleted when either the
parent membership or the parent resource is removed.
@@ -1251,18 +1253,19 @@ Create a new override.
{
"compute_allocation_membership_id": "membership-123",
"compute_allocation_resource_id": "resource-456",
- "overridden_resource_amount": 10000
+ "override_resource_amount": 24,
+ "override_resource_time": 1440
}
```
-- All three fields are required.
-- `overridden_resource_amount` must be non-negative.
+- `compute_allocation_membership_id` and `compute_allocation_resource_id` are required.
+- `override_resource_amount` and `override_resource_time` must be non-negative (default `0`).
- `id` is generated server-side when omitted.
**Errors**
- `400` — missing required fields, referenced membership/resource not found,
- or negative `overridden_resource_amount`.
+ or negative `override_resource_amount` / `override_resource_time`.
- `409` — an override already exists for this `(membership, resource)` pair.
### GET `/compute-allocation-membership-resource-overrides/{id}`
@@ -1367,19 +1370,20 @@ curl -s $BASE/compute-allocations/$ALLOC_ID/memberships | jq
# List every allocation this user is a member of.
curl -s $BASE/users/$CLUSTER_USER_ACCT_ID/compute-allocation-memberships | jq
-# Grant the membership a per-resource SU override on the debug resource.
+# Grant the membership a per-resource override on the debug resource.
OVERRIDE_ID=$(curl -s -X POST $BASE/compute-allocation-membership-resource-overrides \
-H 'Content-Type: application/json' \
-d "{
\"compute_allocation_membership_id\":\"$MEMBERSHIP_ID\",
\"compute_allocation_resource_id\":\"$RES_ID\",
- \"overridden_resource_amount\":10
+ \"override_resource_amount\":1,
+ \"override_resource_time\":20
}" | jq -r .id)
-# Bump the override amount.
+# Bump the override amount and time.
curl -s -X PUT $BASE/compute-allocation-membership-resource-overrides/$OVERRIDE_ID \
-H 'Content-Type: application/json' \
- -d '{"overridden_resource_amount":15000}' | jq
+ -d '{"override_resource_amount":16,"override_resource_time":960}' | jq
# List every resource override for the membership.
curl -s $BASE/compute-allocation-memberships/$MEMBERSHIP_ID/resource-overrides | jq
diff --git a/internal/db/migrations/000013_compute_allocation_membership_resource_override_amounts.down.sql b/internal/db/migrations/000013_compute_allocation_membership_resource_override_amounts.down.sql
new file mode 100644
index 000000000..c937d1aca
--- /dev/null
+++ b/internal/db/migrations/000013_compute_allocation_membership_resource_override_amounts.down.sql
@@ -0,0 +1,20 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+ALTER TABLE compute_allocation_membership_resource_overrides
+ DROP COLUMN override_resource_time,
+ CHANGE COLUMN override_resource_amount overridden_resource_amount BIGINT NOT NULL DEFAULT 0;
diff --git a/internal/db/migrations/000013_compute_allocation_membership_resource_override_amounts.up.sql b/internal/db/migrations/000013_compute_allocation_membership_resource_override_amounts.up.sql
new file mode 100644
index 000000000..bc18ef857
--- /dev/null
+++ b/internal/db/migrations/000013_compute_allocation_membership_resource_override_amounts.up.sql
@@ -0,0 +1,20 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+ALTER TABLE compute_allocation_membership_resource_overrides
+ CHANGE COLUMN overridden_resource_amount override_resource_amount BIGINT NOT NULL DEFAULT 0,
+ ADD COLUMN override_resource_time BIGINT NOT NULL DEFAULT 0 AFTER override_resource_amount;
diff --git a/internal/store/compute_allocation_membership_resource_override_store.go b/internal/store/compute_allocation_membership_resource_override_store.go
index 1e0f0d1b0..d89f11e9c 100644
--- a/internal/store/compute_allocation_membership_resource_override_store.go
+++ b/internal/store/compute_allocation_membership_resource_override_store.go
@@ -27,7 +27,7 @@ import (
"github.com/apache/airavata-custos/pkg/models"
)
-const computeAllocationMembershipResourceOverrideColumns = "id, compute_allocation_membership_id, compute_allocation_resource_id, overridden_resource_amount"
+const computeAllocationMembershipResourceOverrideColumns = "id, compute_allocation_membership_id, compute_allocation_resource_id, override_resource_amount, override_resource_time"
type mysqlComputeAllocationMembershipResourceOverrideStore struct {
db *sqlx.DB
@@ -99,9 +99,9 @@ func (s *mysqlComputeAllocationMembershipResourceOverrideStore) FindByResource(c
func (s *mysqlComputeAllocationMembershipResourceOverrideStore) Create(ctx context.Context, tx *sql.Tx, o *models.ComputeAllocationMembershipResourceOverride) error {
_, err := tx.ExecContext(ctx,
`INSERT INTO compute_allocation_membership_resource_overrides
- (id, compute_allocation_membership_id, compute_allocation_resource_id, overridden_resource_amount)
- VALUES (?, ?, ?, ?)`,
- o.ID, o.ComputeAllocationMembershipID, o.ComputeAllocationResourceID, o.OverriddenResourceAmount)
+ (id, compute_allocation_membership_id, compute_allocation_resource_id, override_resource_amount, override_resource_time)
+ VALUES (?, ?, ?, ?, ?)`,
+ o.ID, o.ComputeAllocationMembershipID, o.ComputeAllocationResourceID, o.OverrideResourceAmount, o.OverrideResourceTime)
return err
}
@@ -110,9 +110,10 @@ func (s *mysqlComputeAllocationMembershipResourceOverrideStore) Update(ctx conte
`UPDATE compute_allocation_membership_resource_overrides
SET compute_allocation_membership_id = ?,
compute_allocation_resource_id = ?,
- overridden_resource_amount = ?
+ override_resource_amount = ?,
+ override_resource_time = ?
WHERE id = ?`,
- o.ComputeAllocationMembershipID, o.ComputeAllocationResourceID, o.OverriddenResourceAmount, o.ID)
+ o.ComputeAllocationMembershipID, o.ComputeAllocationResourceID, o.OverrideResourceAmount, o.OverrideResourceTime, o.ID)
return err
}
diff --git a/pkg/models/allocation.go b/pkg/models/allocation.go
index 582bce606..a8fcd3976 100644
--- a/pkg/models/allocation.go
+++ b/pkg/models/allocation.go
@@ -103,7 +103,8 @@ type ComputeAllocationMembershipResourceOverride struct {
ID string `json:"id" db:"id"`
ComputeAllocationMembershipID string `json:"compute_allocation_membership_id" db:"compute_allocation_membership_id"`
ComputeAllocationResourceID string `json:"compute_allocation_resource_id" db:"compute_allocation_resource_id"`
- OverriddenResourceAmount int64 `json:"overridden_resource_amount" db:"overridden_resource_amount"` // The overridden SU amount for the user for this specific resource, e.g., 150 SUs for GPU hours, etc.
+ OverrideResourceAmount int64 `json:"override_resource_amount" db:"override_resource_amount"` // The overridden amount of the resource for the user (e.g., number of CPUs, GPUs).
+ OverrideResourceTime int64 `json:"override_resource_time" db:"override_resource_time"` // The overridden wall-clock time in minutes that the resource amount is granted for.
}
type ComputeAllocationMembership struct {
diff --git a/pkg/service/compute_allocation_membership_resource_override.go b/pkg/service/compute_allocation_membership_resource_override.go
index 81fadee40..e3148cd00 100644
--- a/pkg/service/compute_allocation_membership_resource_override.go
+++ b/pkg/service/compute_allocation_membership_resource_override.go
@@ -40,8 +40,11 @@ func (s *Service) CreateComputeAllocationMembershipResourceOverride(ctx context.
if o.ComputeAllocationResourceID == "" {
return nil, fmt.Errorf("%w: compute_allocation_resource_id is required", ErrInvalidInput)
}
- if o.OverriddenResourceAmount < 0 {
- return nil, fmt.Errorf("%w: overridden_resource_amount must be non-negative", ErrInvalidInput)
+ if o.OverrideResourceAmount < 0 {
+ return nil, fmt.Errorf("%w: override_resource_amount must be non-negative", ErrInvalidInput)
+ }
+ if o.OverrideResourceTime < 0 {
+ return nil, fmt.Errorf("%w: override_resource_time must be non-negative", ErrInvalidInput)
}
if m, err := s.memberships.FindByID(ctx, o.ComputeAllocationMembershipID); err != nil {
@@ -141,8 +144,11 @@ func (s *Service) UpdateComputeAllocationMembershipResourceOverride(ctx context.
if o == nil || o.ID == "" {
return nil, fmt.Errorf("%w: membership resource override id is required", ErrInvalidInput)
}
- if o.OverriddenResourceAmount < 0 {
- return nil, fmt.Errorf("%w: overridden_resource_amount must be non-negative", ErrInvalidInput)
+ if o.OverrideResourceAmount < 0 {
+ return nil, fmt.Errorf("%w: override_resource_amount must be non-negative", ErrInvalidInput)
+ }
+ if o.OverrideResourceTime < 0 {
+ return nil, fmt.Errorf("%w: override_resource_time must be non-negative", ErrInvalidInput)
}
existing, err := s.membershipOverrides.FindByID(ctx, o.ID)
if err != nil {
(airavata-custos) 02/07: Creating a slurm account when an compute allocation was created
Posted by di...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
DImuthuUpe pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/airavata-custos.git
commit 983beccd15dba842fa601258e6d85be23039f4f2
Author: DImuthuUpe <di...@gmail.com>
AuthorDate: Mon May 18 02:11:10 2026 -0400
Creating a slurm account when an compute allocation was created
---
cmd/server/main.go | 4 +-
connectors/SLURM/Association-Mapper/README.md | 30 +++++++++++++
.../internal/operations/accounts.go | 8 ++--
.../internal/operations/accounts_test.go | 8 ++--
.../internal/operations/associations.go | 6 +--
.../internal/operations/associations_test.go | 7 ++--
.../internal/operations/client.go | 20 +++++----
.../internal/operations/client_test.go | 4 +-
.../internal/subscribers/account.go | 49 +++++++++++++++++++++-
.../internal/subscribers/subscriber.go | 5 ++-
.../SLURM/Association-Mapper/pkg/smapper/loader.go | 22 ++++++++--
docs/API-Docs.md | 4 +-
internal/connectors/loader.go | 5 ++-
13 files changed, 134 insertions(+), 38 deletions(-)
diff --git a/cmd/server/main.go b/cmd/server/main.go
index ea0ebd697..e2cc07006 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -72,12 +72,12 @@ func run() error {
// Create a new event bus instance to async messaging between service and connectors
eventBus := events.New()
+ svc := service.New(database, eventBus)
- if err := connectors.LoadConnectors(eventBus); err != nil {
+ if err := connectors.LoadConnectors(eventBus, svc); err != nil {
return err
}
- svc := service.New(database, eventBus)
handler := server.LoggingMiddleware(server.New(svc))
httpServer := &http.Server{
diff --git a/connectors/SLURM/Association-Mapper/README.md b/connectors/SLURM/Association-Mapper/README.md
index 9d9bfc6f0..2dff57ba3 100644
--- a/connectors/SLURM/Association-Mapper/README.md
+++ b/connectors/SLURM/Association-Mapper/README.md
@@ -19,6 +19,36 @@ This package is part of the root `github.com/apache/airavata-custos` module.
Import path: `github.com/apache/airavata-custos/connectors/SLURM/Association-Mapper/internal/operations`.
+## Configuration
+
+The connector reads its `slurmrestd` connection details from the following environment variables:
+
+| Variable | Description |
+|---------------|-----------------------------------------------------------------------------|
+| `SLURM_API` | Base URL of the `slurmrestd` endpoint, e.g. `https://slurm.example.org:6820` |
+| `SLURM_USER` | SLURM user name to authenticate as (sent in the `X-SLURM-USER-NAME` header) |
+| `SLURM_TOKEN` | JWT token for that user (sent in the `X-SLURM-USER-TOKEN` header) |
+| `SLURM_API_VERSION` | API version of `slurmrestd`. You can check it from `slurmrestd -d list` |
+
+Example:
+
+```bash
+export SLURM_API='https://slurm.example.org:6820'
+export SLURM_USER='slurm'
+export SLURM_TOKEN="$(scontrol token lifespan=3600 | awk -F= '{print $2}')"
+export SLURM_API_VERSION='40'
+```
+
+All three are required for any live `slurmrestd` interaction. Tests do not need them.
+
+You can check the functionality of the SLURM REST API through
+```bash
+curl -sS \
+ -H "X-SLURM-USER-NAME: $SLURM_USER" \
+ -H "X-SLURM-USER-TOKEN: $SLURM_TOKEN" \
+ "$SLURM_API/slurm/v0.0.$SLURM_API_VERSION/ping" | jq
+```
+
## Test
```bash
diff --git a/connectors/SLURM/Association-Mapper/internal/operations/accounts.go b/connectors/SLURM/Association-Mapper/internal/operations/accounts.go
index ea7d87708..18cc0e456 100644
--- a/connectors/SLURM/Association-Mapper/internal/operations/accounts.go
+++ b/connectors/SLURM/Association-Mapper/internal/operations/accounts.go
@@ -9,7 +9,7 @@ type accountsResponse struct {
func (c *Client) ListAccounts() ([]Account, error) {
var out accountsResponse
- if _, err := c.do("GET", "/slurmdb/v0.0.41/accounts", nil, &out); err != nil {
+ if _, err := c.do("GET", "/slurmdb/v0.0."+c.apiVersion+"/accounts", nil, &out); err != nil {
return nil, err
}
return out.Accounts, nil
@@ -17,7 +17,7 @@ func (c *Client) ListAccounts() ([]Account, error) {
func (c *Client) GetAccount(name string) (*Account, error) {
var out accountsResponse
- if _, err := c.do("GET", "/slurmdb/v0.0.41/account/"+name, nil, &out); err != nil {
+ if _, err := c.do("GET", "/slurmdb/v0.0."+c.apiVersion+"/account/"+name, nil, &out); err != nil {
return nil, err
}
if len(out.Accounts) == 0 {
@@ -43,11 +43,11 @@ func (c *Client) CreateAccount(a Account, cluster string) error {
"organization": a.Organization,
},
}
- _, err := c.do("POST", "/slurmdb/v0.0.41/accounts_association/", body, nil)
+ _, err := c.do("POST", "/slurmdb/v0.0."+c.apiVersion+"/accounts_association/", body, nil)
return err
}
func (c *Client) DeleteAccount(name string) error {
- _, err := c.do("DELETE", "/slurmdb/v0.0.41/account/"+name, nil, nil)
+ _, err := c.do("DELETE", "/slurmdb/v0.0."+c.apiVersion+"/account/"+name, nil, nil)
return err
}
diff --git a/connectors/SLURM/Association-Mapper/internal/operations/accounts_test.go b/connectors/SLURM/Association-Mapper/internal/operations/accounts_test.go
index a427116e5..aaeae99df 100644
--- a/connectors/SLURM/Association-Mapper/internal/operations/accounts_test.go
+++ b/connectors/SLURM/Association-Mapper/internal/operations/accounts_test.go
@@ -18,7 +18,7 @@ func TestListAccounts(t *testing.T) {
}))
defer srv.Close()
- c := New(srv.URL, "root", "t")
+ c := New(srv.URL, "root", "t", "41")
accts, err := c.ListAccounts()
if err != nil {
t.Fatal(err)
@@ -58,7 +58,7 @@ func TestCreateAccount(t *testing.T) {
}))
defer srv.Close()
- c := New(srv.URL, "root", "t")
+ c := New(srv.URL, "root", "t", "41")
if err := c.CreateAccount(Account{Name: "eng", Description: "engineering"}, "artisan"); err != nil {
t.Fatal(err)
}
@@ -73,7 +73,7 @@ func TestDeleteAccount(t *testing.T) {
}))
defer srv.Close()
- c := New(srv.URL, "root", "t")
+ c := New(srv.URL, "root", "t", "41")
if err := c.DeleteAccount("eng"); err != nil {
t.Fatal(err)
}
@@ -88,7 +88,7 @@ func TestGetAccount(t *testing.T) {
}))
defer srv.Close()
- c := New(srv.URL, "root", "t")
+ c := New(srv.URL, "root", "t", "41")
a, err := c.GetAccount("eng")
if err != nil {
t.Fatal(err)
diff --git a/connectors/SLURM/Association-Mapper/internal/operations/associations.go b/connectors/SLURM/Association-Mapper/internal/operations/associations.go
index ab34cd931..fc1607480 100644
--- a/connectors/SLURM/Association-Mapper/internal/operations/associations.go
+++ b/connectors/SLURM/Association-Mapper/internal/operations/associations.go
@@ -32,7 +32,7 @@ type associationsResponse struct {
}
func (c *Client) ListAssociations(f AssocFilter) ([]Association, error) {
- path := "/slurmdb/v0.0.41/associations"
+ path := "/slurmdb/v0.0." + c.apiVersion + "/associations"
if q := f.query(); q != "" {
path += "?" + q
}
@@ -48,12 +48,12 @@ func (c *Client) ListAssociations(f AssocFilter) ([]Association, error) {
// triple exists, it's updated; otherwise created.
func (c *Client) UpsertAssociation(a Association) error {
body := map[string]any{"associations": []Association{a}}
- _, err := c.do("POST", "/slurmdb/v0.0.41/associations", body, nil)
+ _, err := c.do("POST", "/slurmdb/v0.0."+c.apiVersion+"/associations", body, nil)
return err
}
func (c *Client) DeleteAssociation(f AssocFilter) error {
- path := "/slurmdb/v0.0.41/association"
+ path := "/slurmdb/v0.0." + c.apiVersion + "/association"
if q := f.query(); q != "" {
path += "?" + q
}
diff --git a/connectors/SLURM/Association-Mapper/internal/operations/associations_test.go b/connectors/SLURM/Association-Mapper/internal/operations/associations_test.go
index 9082927b0..d16da5186 100644
--- a/connectors/SLURM/Association-Mapper/internal/operations/associations_test.go
+++ b/connectors/SLURM/Association-Mapper/internal/operations/associations_test.go
@@ -10,7 +10,6 @@ import (
"testing"
)
-
func TestListAssociationsByAccount(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/slurmdb/v0.0.41/associations" {
@@ -22,7 +21,7 @@ func TestListAssociationsByAccount(t *testing.T) {
_, _ = w.Write([]byte(`{"associations":[{"account":"eng","cluster":"artisan","user":"alice","id_association":5}]}`))
}))
defer srv.Close()
- c := New(srv.URL, "root", "t")
+ c := New(srv.URL, "root", "t", "41")
assocs, err := c.ListAssociations(AssocFilter{Account: "eng"})
if err != nil {
t.Fatal(err)
@@ -49,7 +48,7 @@ func TestCreateAssociation(t *testing.T) {
_, _ = w.Write([]byte(`{}`))
}))
defer srv.Close()
- c := New(srv.URL, "root", "t")
+ c := New(srv.URL, "root", "t", "41")
err := c.UpsertAssociation(Association{Account: "eng", Cluster: "artisan", User: "alice"})
if err != nil {
t.Fatal(err)
@@ -67,7 +66,7 @@ func TestDeleteAssociation(t *testing.T) {
_, _ = w.Write([]byte(`{}`))
}))
defer srv.Close()
- c := New(srv.URL, "root", "t")
+ c := New(srv.URL, "root", "t", "41")
if err := c.DeleteAssociation(AssocFilter{Account: "eng", User: "alice"}); err != nil {
t.Fatal(err)
}
diff --git a/connectors/SLURM/Association-Mapper/internal/operations/client.go b/connectors/SLURM/Association-Mapper/internal/operations/client.go
index 62bd25b99..a60ab883a 100644
--- a/connectors/SLURM/Association-Mapper/internal/operations/client.go
+++ b/connectors/SLURM/Association-Mapper/internal/operations/client.go
@@ -12,18 +12,20 @@ import (
)
type Client struct {
- baseURL string
- user string
- token string
- http *http.Client
+ baseURL string
+ user string
+ token string
+ apiVersion string
+ http *http.Client
}
-func New(baseURL, user, token string) *Client {
+func New(baseURL, user, token, apiVersion string) *Client {
return &Client{
- baseURL: strings.TrimRight(baseURL, "/"),
- user: user,
- token: token,
- http: &http.Client{Timeout: 30 * time.Second},
+ baseURL: strings.TrimRight(baseURL, "/"),
+ user: user,
+ token: token,
+ apiVersion: apiVersion,
+ http: &http.Client{Timeout: 30 * time.Second},
}
}
diff --git a/connectors/SLURM/Association-Mapper/internal/operations/client_test.go b/connectors/SLURM/Association-Mapper/internal/operations/client_test.go
index 3d07b6456..c73d3ecb5 100644
--- a/connectors/SLURM/Association-Mapper/internal/operations/client_test.go
+++ b/connectors/SLURM/Association-Mapper/internal/operations/client_test.go
@@ -18,7 +18,7 @@ func TestDoSetsHeaders(t *testing.T) {
}))
defer srv.Close()
- c := New(srv.URL, "root", "tok123")
+ c := New(srv.URL, "root", "tok123", "41")
if _, err := c.do("GET", "/slurm/v0.0.41/ping", nil, nil); err != nil {
t.Fatal(err)
}
@@ -36,7 +36,7 @@ func TestDoUnwrapsErrors(t *testing.T) {
_, _ = w.Write([]byte(`{"errors":[{"description":"nope","error_number":42,"error":"Bad","source":"test"}]}`))
}))
defer srv.Close()
- c := New(srv.URL, "root", "tok")
+ c := New(srv.URL, "root", "tok", "41")
_, err := c.do("GET", "/x", nil, nil)
if err == nil {
t.Fatal("expected error")
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/account.go b/connectors/SLURM/Association-Mapper/internal/subscribers/account.go
index 9c2f29c6f..a4a41f99c 100644
--- a/connectors/SLURM/Association-Mapper/internal/subscribers/account.go
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/account.go
@@ -1,10 +1,55 @@
package subscribers
-import "github.com/apache/airavata-custos/pkg/models"
-import "log/slog"
+import (
+ "log/slog"
+
+ "context"
+ "time"
+
+ client "github.com/apache/airavata-custos/connectors/SLURM/Association-Mapper/internal/operations"
+ "github.com/apache/airavata-custos/pkg/models"
+)
func (a *AssociationSubscriber) SubscribeToComputeAccountCreation(computeAccount models.ComputeAllocation) {
slog.Info("Received compute account creation event", "account", computeAccount)
+
+ ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+ defer cancel()
+
+ cluster, err := a.coreService.GetComputeCluster(ctx, computeAccount.ComputeClusterID)
+ if err != nil {
+ slog.Error("Failed to get compute cluster for account creation", "error", err)
+ return
+ }
+
+ project, err := a.coreService.GetProject(ctx, computeAccount.ProjectID)
+ if err != nil {
+ slog.Error("Failed to get project for account creation", "error", err)
+ return
+ }
+
+ pi, err := a.coreService.GetUser(ctx, project.ProjectPIID)
+ if err != nil {
+ slog.Error("Failed to get project PI for account creation", "error", err)
+ return
+ }
+
+ organization, err := a.coreService.GetOrganization(ctx, pi.OrganizationID)
+ if err != nil {
+ slog.Error("Failed to get organization for account creation", "error", err)
+ return
+ }
+
+ slurmAccount := client.Account{
+ Name: computeAccount.Name,
+ Description: computeAccount.Name,
+ Organization: organization.Name,
+ }
+
+ err = a.slurmClient.CreateAccount(slurmAccount, cluster.Name) // TODO: where to get cluster name from?
+ if err != nil {
+ slog.Error("Failed to create SLURM account", "error", err)
+ }
}
func (a *AssociationSubscriber) SubscribeToComputeAccountDeletion(computeAccount models.ComputeAllocation) {
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go b/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go
index 129bfcff9..30b03f9af 100644
--- a/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go
@@ -2,16 +2,19 @@ package subscribers
import client "github.com/apache/airavata-custos/connectors/SLURM/Association-Mapper/internal/operations"
import "github.com/apache/airavata-custos/pkg/events"
+import "github.com/apache/airavata-custos/pkg/service"
type AssociationSubscriber struct {
slurmClient *client.Client
eventBus *events.Bus
+ coreService *service.Service
}
-func NewAssociationSubscriber(slurmClient *client.Client, eventBus *events.Bus) *AssociationSubscriber {
+func NewAssociationSubscriber(slurmClient *client.Client, eventBus *events.Bus, coreService *service.Service) *AssociationSubscriber {
return &AssociationSubscriber{
slurmClient: slurmClient,
eventBus: eventBus,
+ coreService: coreService,
}
}
diff --git a/connectors/SLURM/Association-Mapper/pkg/smapper/loader.go b/connectors/SLURM/Association-Mapper/pkg/smapper/loader.go
index 9c172b4f8..8f5dc39be 100644
--- a/connectors/SLURM/Association-Mapper/pkg/smapper/loader.go
+++ b/connectors/SLURM/Association-Mapper/pkg/smapper/loader.go
@@ -3,9 +3,25 @@ package smapper
import "github.com/apache/airavata-custos/pkg/events"
import "github.com/apache/airavata-custos/connectors/SLURM/Association-Mapper/internal/subscribers"
import client "github.com/apache/airavata-custos/connectors/SLURM/Association-Mapper/internal/operations"
+import "os"
+import "log/slog"
+import "github.com/apache/airavata-custos/pkg/service"
-func LoadConnector(eventBus *events.Bus) error {
- slurmClient := client.New("localhost:8080", "", "") // Replace with actual SLURM client initialization.
- subscribers.NewAssociationSubscriber(slurmClient, eventBus).RegisterSubscribers()
+func LoadConnector(eventBus *events.Bus, coreService *service.Service) error {
+
+ // Read url, username, and password from environment variables
+ apiUrl := os.Getenv("SLURM_API")
+ user := os.Getenv("SLURM_USER")
+ token := os.Getenv("SLURM_TOKEN")
+ apiVersion := os.Getenv("SLURM_API_VERSION")
+ if apiUrl == "" || user == "" || token == "" || apiVersion == "" {
+ slog.Info("SLURM API credentials not fully provided, skipping SLURM Association Mapper connector")
+ // print valualues of the env vars for debugging
+ slog.Info("SLURM API credentials", "apiUrl", apiUrl, "user", user, "token", token, "apiVersion", apiVersion)
+ return nil // skip loading if any of the required env vars are missing
+ }
+
+ slurmClient := client.New(apiUrl, user, token, apiVersion)
+ subscribers.NewAssociationSubscriber(slurmClient, eventBus, coreService).RegisterSubscribers()
return nil
}
diff --git a/docs/API-Docs.md b/docs/API-Docs.md
index 2a308db05..6931d0304 100644
--- a/docs/API-Docs.md
+++ b/docs/API-Docs.md
@@ -1052,7 +1052,7 @@ BASE=http://localhost:8080
ORG_ID=$(curl -s -X POST $BASE/organizations \
-H 'Content-Type: application/json' \
- -d '{"name":"University of Example","originated_id":"ACCESS-ORG-001"}' \
+ -d '{"name":"Georgia Institute of Technology","originated_id":"ACCESS-ORG-001"}' \
| jq -r .id)
USER_ID=$(curl -s -X POST $BASE/users \
@@ -1067,7 +1067,7 @@ PROJ_ID=$(curl -s -X POST $BASE/projects \
CLUSTER_ID=$(curl -s -X POST $BASE/compute-clusters \
-H 'Content-Type: application/json' \
- -d '{"name":"Delta"}' | jq -r .id)
+ -d '{"name":"nexus-dev"}' | jq -r .id)
ALLOC_ID=$(curl -s -X POST $BASE/compute-allocations \
-H 'Content-Type: application/json' \
diff --git a/internal/connectors/loader.go b/internal/connectors/loader.go
index 1b603abd1..c1e129661 100644
--- a/internal/connectors/loader.go
+++ b/internal/connectors/loader.go
@@ -3,13 +3,14 @@ package connectors
import "github.com/apache/airavata-custos/connectors/SLURM/Association-Mapper/pkg/smapper"
import "github.com/apache/airavata-custos/pkg/events"
import "log/slog"
+import "github.com/apache/airavata-custos/pkg/service"
-func LoadConnectors(eventBus *events.Bus) error {
+func LoadConnectors(eventBus *events.Bus, coreService *service.Service) error {
slog.Info("loading connectors")
slog.Info("loading SLURM Association Mapper connector")
- err := smapper.LoadConnector(eventBus)
+ err := smapper.LoadConnector(eventBus, coreService)
if err != nil {
slog.Error("failed to load SLURM Association Mapper connector", "error", err)
return err
(airavata-custos) 03/07: Fixing minor misconfiguration
Posted by di...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
DImuthuUpe pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/airavata-custos.git
commit 9b7c6e5440bd12b4df12c3a3e4e94f865f08a181
Author: DImuthuUpe <di...@gmail.com>
AuthorDate: Mon May 18 08:56:52 2026 -0400
Fixing minor misconfiguration
---
.../internal/subscribers/account.go | 20 ++++++++++----------
.../internal/subscribers/subscriber.go | 6 +++---
pkg/service/compute_cluster.go | 6 +++---
3 files changed, 16 insertions(+), 16 deletions(-)
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/account.go b/connectors/SLURM/Association-Mapper/internal/subscribers/account.go
index a4a41f99c..2b872d54c 100644
--- a/connectors/SLURM/Association-Mapper/internal/subscribers/account.go
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/account.go
@@ -10,33 +10,33 @@ import (
"github.com/apache/airavata-custos/pkg/models"
)
-func (a *AssociationSubscriber) SubscribeToComputeAccountCreation(computeAccount models.ComputeAllocation) {
- slog.Info("Received compute account creation event", "account", computeAccount)
+func (a *AssociationSubscriber) SubscribeToComputeAllocationCreation(computeAccount models.ComputeAllocation) {
+ slog.Info("Received compute allocation creation event", "account", computeAccount)
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
cluster, err := a.coreService.GetComputeCluster(ctx, computeAccount.ComputeClusterID)
if err != nil {
- slog.Error("Failed to get compute cluster for account creation", "error", err)
+ slog.Error("Failed to get compute cluster for allocation creation", "error", err)
return
}
project, err := a.coreService.GetProject(ctx, computeAccount.ProjectID)
if err != nil {
- slog.Error("Failed to get project for account creation", "error", err)
+ slog.Error("Failed to get project for allocation creation", "error", err)
return
}
pi, err := a.coreService.GetUser(ctx, project.ProjectPIID)
if err != nil {
- slog.Error("Failed to get project PI for account creation", "error", err)
+ slog.Error("Failed to get project PI for allocation creation", "error", err)
return
}
organization, err := a.coreService.GetOrganization(ctx, pi.OrganizationID)
if err != nil {
- slog.Error("Failed to get organization for account creation", "error", err)
+ slog.Error("Failed to get organization for allocation creation", "error", err)
return
}
@@ -52,10 +52,10 @@ func (a *AssociationSubscriber) SubscribeToComputeAccountCreation(computeAccount
}
}
-func (a *AssociationSubscriber) SubscribeToComputeAccountDeletion(computeAccount models.ComputeAllocation) {
- slog.Info("Received compute account deletion event", "account", computeAccount)
+func (a *AssociationSubscriber) SubscribeToComputeAllocationDeletion(computeAccount models.ComputeAllocation) {
+ slog.Info("Received compute allocation deletion event", "account", computeAccount)
}
-func (a *AssociationSubscriber) SubscribeToComputeAccountUpdate(computeAccount models.ComputeAllocation) {
- slog.Info("Received compute account update event", "account", computeAccount)
+func (a *AssociationSubscriber) SubscribeToComputeAllocationUpdate(computeAccount models.ComputeAllocation) {
+ slog.Info("Received compute allocation update event", "account", computeAccount)
}
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go b/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go
index 30b03f9af..ce5161fd7 100644
--- a/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go
@@ -19,7 +19,7 @@ func NewAssociationSubscriber(slurmClient *client.Client, eventBus *events.Bus,
}
func (a *AssociationSubscriber) RegisterSubscribers() {
- a.eventBus.SubscribeComputeAllocationCreated(a.SubscribeToComputeAccountCreation)
- a.eventBus.SubscribeComputeAllocationDeleted(a.SubscribeToComputeAccountDeletion)
- a.eventBus.SubscribeComputeAllocationUpdated(a.SubscribeToComputeAccountUpdate)
+ a.eventBus.SubscribeComputeAllocationCreated(a.SubscribeToComputeAllocationCreation)
+ a.eventBus.SubscribeComputeAllocationDeleted(a.SubscribeToComputeAllocationDeletion)
+ a.eventBus.SubscribeComputeAllocationUpdated(a.SubscribeToComputeAllocationUpdate)
}
diff --git a/pkg/service/compute_cluster.go b/pkg/service/compute_cluster.go
index 41198709b..a9d5b3091 100644
--- a/pkg/service/compute_cluster.go
+++ b/pkg/service/compute_cluster.go
@@ -51,7 +51,7 @@ func (s *Service) CreateComputeCluster(ctx context.Context, cluster *models.Comp
return nil, fmt.Errorf("create compute cluster: %w", err)
}
- s.eventBus.Publish(events.ComputeAllocationCreateEvent, cluster)
+ s.eventBus.Publish(events.ComputeClusterCreateEvent, cluster)
return cluster, nil
}
@@ -103,7 +103,7 @@ func (s *Service) UpdateComputeCluster(ctx context.Context, cluster *models.Comp
return fmt.Errorf("update compute cluster: %w", err)
}
- s.eventBus.Publish(events.ComputeAllocationUpdateEvent, cluster)
+ s.eventBus.Publish(events.ComputeClusterUpdateEvent, cluster)
return nil
}
@@ -125,6 +125,6 @@ func (s *Service) DeleteComputeCluster(ctx context.Context, id string) error {
return fmt.Errorf("delete compute cluster: %w", err)
}
- s.eventBus.Publish(events.ComputeAllocationDeleteEvent, cluster)
+ s.eventBus.Publish(events.ComputeClusterDeleteEvent, cluster)
return nil
}
(airavata-custos) 05/07: Moving member-specific quota to a seperate data model and implementing slurm subscriber
Posted by di...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
DImuthuUpe pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/airavata-custos.git
commit 3b0bf6845ed2b1f949eba306180e2c2ae1b6db06
Author: DImuthuUpe <di...@gmail.com>
AuthorDate: Tue May 19 14:36:11 2026 -0400
Moving member-specific quota to a seperate data model and implementing slurm subscriber
---
.../internal/subscribers/members.go | 94 +++++++++-
.../internal/subscribers/subscriber.go | 1 +
docs/API-Docs.md | 94 +++++++---
...location_membership_resource_overrides.down.sql | 21 +++
...allocation_membership_resource_overrides.up.sql | 36 ++++
internal/server/server.go | 88 ++++++++--
...llocation_membership_resource_override_store.go | 123 +++++++++++++
.../store/compute_allocation_membership_store.go | 11 +-
internal/store/store.go | 20 +++
...ation_membership_resource_override_subscribe.go | 67 ++++++++
pkg/events/types.go | 7 +
pkg/models/allocation.go | 14 +-
pkg/service/compute_allocation_membership.go | 27 ---
...pute_allocation_membership_resource_override.go | 191 +++++++++++++++++++++
pkg/service/service.go | 4 +
15 files changed, 721 insertions(+), 77 deletions(-)
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/members.go b/connectors/SLURM/Association-Mapper/internal/subscribers/members.go
index cd609d5af..687ecd8d5 100644
--- a/connectors/SLURM/Association-Mapper/internal/subscribers/members.go
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/members.go
@@ -4,12 +4,15 @@ import (
"log/slog"
"context"
+ "time"
+
client "github.com/apache/airavata-custos/connectors/SLURM/Association-Mapper/internal/operations"
"github.com/apache/airavata-custos/pkg/models"
- "time"
)
-func (a *AssociationSubscriber) SubscribeToComputeAllocationMembershipCreation(membership models.ComputeAllocationMembership) {
+func (a *AssociationSubscriber) SubscribeToComputeAllocationMembershipCreation(
+ membership models.ComputeAllocationMembership) {
+
slog.Info("Received compute allocation membership creation event", "membership", membership)
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
@@ -48,6 +51,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationMembershipCreation(m
if len(resources) == 0 {
slog.Warn("No resources found for allocation, defaulting to partition 'default'", "allocation_id", allocation.ID)
}
+
association := client.Association{
Account: allocation.Name,
Cluster: cluster.Name,
@@ -62,3 +66,89 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationMembershipCreation(m
slog.Info("Successfully upserted association", "association", association)
}
}
+
+func (a *AssociationSubscriber) SubscribeToComputeAllocationMembershipResourceOverrideCreation(
+ override models.ComputeAllocationMembershipResourceOverride) {
+
+ slog.Info("Received compute allocation membership resource override creation event", "override", override)
+
+ ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+ defer cancel()
+ // TODO: read per-resource override via
+ membership, err := a.coreService.GetComputeAllocationMembership(ctx, override.ComputeAllocationMembershipID)
+ if err != nil {
+ slog.Error("Failed to get compute allocation membership for resource override creation", "error", err)
+ return
+ }
+
+ allocationResource, err := a.coreService.GetComputeAllocationResource(ctx, override.ComputeAllocationResourceID)
+ if err != nil {
+ slog.Error("Failed to get compute allocation resource for resource override creation", "error", err)
+ return
+ }
+
+ slog.Info("Received compute allocation membership resource override creation event",
+ "membership", membership, "resource", allocationResource, "override", override)
+
+ allocation, err := a.coreService.GetComputeAllocation(ctx, membership.ComputeAllocationID)
+ if err != nil {
+ slog.Error("Failed to get compute allocation", "error", err)
+ return
+ }
+
+ cluster, err := a.coreService.GetComputeCluster(ctx, allocation.ComputeClusterID)
+ if err != nil {
+ slog.Error("Failed to get compute cluster", "error", err)
+ return
+ }
+
+ user, err := a.coreService.GetUser(ctx, membership.UserID)
+ if err != nil {
+ slog.Error("Failed to get user", "error", err)
+ return
+ }
+
+ csu, err := a.coreService.GetComputeClusterUserByPair(ctx, cluster.ID, user.ID) // TODO: use this to get the local username for the association instead of assuming it's the same as the Airavata Custos username
+ if err != nil {
+ slog.Error("Failed to get compute cluster user by pair", "error", err)
+ return
+ }
+
+ grpTres := []client.TRES{}
+
+ if allocationResource.ResourceType == "GrpTRES" {
+ grpTres = append(grpTres, client.TRES{
+ Type: allocationResource.Name,
+ Count: override.OverriddenResourceAmount, // override.OverriddenResourceAmount is the SU amount, but SLURM needs the actual resource amount (e.g., number of CPU hours), so we need to convert it using the rate for the resource
+ })
+ }
+
+ grpTresMins := []client.TRES{}
+
+ if allocationResource.ResourceType == "GrpTRESMins" {
+ grpTresMins = append(grpTresMins, client.TRES{
+ Type: allocationResource.Name,
+ Count: override.OverriddenResourceAmount,
+ })
+ }
+
+ limits := client.AssocLimits{
+ GrpTRES: grpTres,
+ GrpTRESMins: grpTresMins,
+ }
+
+ association := client.Association{
+ Account: allocation.Name,
+ Cluster: cluster.Name,
+ User: csu.LocalUsername,
+ Partition: allocationResource.Name,
+ Limits: limits,
+ }
+
+ err = a.slurmClient.UpsertAssociation(association)
+ if err != nil {
+ slog.Error("Failed to upsert association for membership resource override creation", "error", err)
+ } else {
+ slog.Info("Successfully upserted association for membership resource override creation", "association", association)
+ }
+}
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go b/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go
index 658dcc073..17b65cce8 100644
--- a/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go
@@ -23,4 +23,5 @@ func (a *AssociationSubscriber) RegisterSubscribers() {
a.eventBus.SubscribeComputeAllocationDeleted(a.SubscribeToComputeAllocationDeletion)
a.eventBus.SubscribeComputeAllocationUpdated(a.SubscribeToComputeAllocationUpdate)
a.eventBus.SubscribeComputeAllocationMembershipCreated(a.SubscribeToComputeAllocationMembershipCreation)
+ a.eventBus.SubscribeComputeAllocationMembershipResourceOverrideCreated(a.SubscribeToComputeAllocationMembershipResourceOverrideCreation)
}
diff --git a/docs/API-Docs.md b/docs/API-Docs.md
index 81aa716db..b571c8255 100644
--- a/docs/API-Docs.md
+++ b/docs/API-Docs.md
@@ -1125,7 +1125,6 @@ Create a new membership.
{
"compute_allocation_id": "alloc-123",
"user_id": "user-456",
- "allocation_amount": 50000,
"start_time": "2026-01-01T00:00:00Z",
"end_time": "2026-12-31T23:59:59Z",
"membership_status": "ACTIVE"
@@ -1136,6 +1135,8 @@ Create a new membership.
existing rows.
- `membership_status` defaults to `ACTIVE` when omitted.
- `id` is generated server-side when omitted.
+- Per-resource SU caps for this membership are stored separately as
+ `ComputeAllocationMembershipResourceOverride` rows; see below.
**Errors**
@@ -1151,21 +1152,6 @@ Retrieve a membership by ID.
Replace mutable fields of a membership. Fields left blank/zero in the request
body fall back to the stored value (partial updates).
-### PUT `/compute-allocation-memberships/{id}/allocation-amount`
-
-Update only the SU sub-allocation granted to the user.
-
-**Request body**
-
-```json
-{ "allocation_amount": 75000 }
-```
-
-**Errors**
-
-- `400` — negative `allocation_amount`.
-- `404` — no membership with the given ID.
-
### PUT `/compute-allocation-memberships/{id}/status`
Update only the lifecycle status of the membership (`ACTIVE`, `INACTIVE`,
@@ -1196,6 +1182,62 @@ List every membership recorded against the given allocation, ordered by
List every allocation membership held by the given user, ordered by
`start_time` ascending.
+### GET `/compute-allocation-memberships/{id}/resource-overrides`
+
+List every per-resource override recorded against the given membership.
+
+---
+
+## Compute Allocation Membership Resource Overrides
+
+A `ComputeAllocationMembershipResourceOverride` records the SU amount of a
+specific resource (`ComputeAllocationResource`) that has been granted to a
+specific membership. There can be at most one override per
+`(compute_allocation_membership_id, compute_allocation_resource_id)` pair
+(enforced by a unique key). Overrides are cascade-deleted when either the
+parent membership or the parent resource is removed.
+
+### POST `/compute-allocation-membership-resource-overrides`
+
+Create a new override.
+
+**Request body**
+
+```json
+{
+ "compute_allocation_membership_id": "membership-123",
+ "compute_allocation_resource_id": "resource-456",
+ "overridden_resource_amount": 10000
+}
+```
+
+- All three fields are required.
+- `overridden_resource_amount` must be non-negative.
+- `id` is generated server-side when omitted.
+
+**Errors**
+
+- `400` — missing required fields, referenced membership/resource not found,
+ or negative `overridden_resource_amount`.
+- `409` — an override already exists for this `(membership, resource)` pair.
+
+### GET `/compute-allocation-membership-resource-overrides/{id}`
+
+Retrieve an override by ID.
+
+### PUT `/compute-allocation-membership-resource-overrides/{id}`
+
+Replace mutable fields of an override. Fields left blank/zero in the request
+body fall back to the stored value (partial updates).
+
+### DELETE `/compute-allocation-membership-resource-overrides/{id}`
+
+Remove an override.
+
+### GET `/compute-allocation-resources/{id}/membership-overrides`
+
+List every membership override referencing the given resource.
+
---
## End-to-end example
@@ -1229,7 +1271,7 @@ ALLOC_ID=$(curl -s -X POST $BASE/compute-allocations \
RES_ID=$(curl -s -X POST $BASE/compute-allocation-resources \
-H 'Content-Type: application/json' \
- -d '{"name":"debug","resource_type":"CPU","resource_amount":24}' | jq -r .id)
+ -d '{"name":"debug","resource_type":"GrpTRES","resource_amount":24}' | jq -r .id)
# Attach the resource to the allocation.
curl -s -X POST $BASE/compute-allocations/$ALLOC_ID/resources \
@@ -1270,7 +1312,6 @@ MEMBERSHIP_ID=$(curl -s -X POST $BASE/compute-allocation-memberships \
-d "{
\"compute_allocation_id\":\"$ALLOC_ID\",
\"user_id\":\"$CLUSTER_USER_ACCT_ID\",
- \"allocation_amount\":50000,
\"start_time\":\"2026-04-01T00:00:00Z\",
\"end_time\":\"2026-06-30T23:59:59Z\",
\"membership_status\":\"ACTIVE\"
@@ -1282,11 +1323,22 @@ curl -s $BASE/compute-allocations/$ALLOC_ID/memberships | jq
# List every allocation this user is a member of.
curl -s $BASE/users/$CLUSTER_USER_ACCT_ID/compute-allocation-memberships | jq
-# Bump the user's SU sub-allocation.
-curl -s -X PUT $BASE/compute-allocation-memberships/$MEMBERSHIP_ID/allocation-amount \
+# Grant the membership a per-resource SU override on the debug resource.
+OVERRIDE_ID=$(curl -s -X POST $BASE/compute-allocation-membership-resource-overrides \
+ -H 'Content-Type: application/json' \
+ -d "{
+ \"compute_allocation_membership_id\":\"$MEMBERSHIP_ID\",
+ \"compute_allocation_resource_id\":\"$RES_ID\",
+ \"overridden_resource_amount\":10000
+ }" | jq -r .id)
+
+# Bump the override amount.
+curl -s -X PUT $BASE/compute-allocation-membership-resource-overrides/$OVERRIDE_ID \
-H 'Content-Type: application/json' \
- -d '{"allocation_amount":75000}' | jq
+ -d '{"overridden_resource_amount":15000}' | jq
+# List every resource override for the membership.
+curl -s $BASE/compute-allocation-memberships/$MEMBERSHIP_ID/resource-overrides | jq
# Record a usage diff against the allocation.
diff --git a/internal/db/migrations/000011_compute_allocation_membership_resource_overrides.down.sql b/internal/db/migrations/000011_compute_allocation_membership_resource_overrides.down.sql
new file mode 100644
index 000000000..e07d3db11
--- /dev/null
+++ b/internal/db/migrations/000011_compute_allocation_membership_resource_overrides.down.sql
@@ -0,0 +1,21 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+DROP TABLE IF EXISTS compute_allocation_membership_resource_overrides;
+
+ALTER TABLE compute_allocation_memberships
+ ADD COLUMN allocation_amount BIGINT NOT NULL DEFAULT 0 AFTER user_id;
diff --git a/internal/db/migrations/000011_compute_allocation_membership_resource_overrides.up.sql b/internal/db/migrations/000011_compute_allocation_membership_resource_overrides.up.sql
new file mode 100644
index 000000000..946bbb029
--- /dev/null
+++ b/internal/db/migrations/000011_compute_allocation_membership_resource_overrides.up.sql
@@ -0,0 +1,36 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+ALTER TABLE compute_allocation_memberships
+ DROP COLUMN allocation_amount;
+
+CREATE TABLE IF NOT EXISTS compute_allocation_membership_resource_overrides
+(
+ id VARCHAR(255) NOT NULL,
+ compute_allocation_membership_id VARCHAR(255) NOT NULL,
+ compute_allocation_resource_id VARCHAR(255) NOT NULL,
+ overridden_resource_amount BIGINT NOT NULL DEFAULT 0,
+ created_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),
+ updated_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6),
+ PRIMARY KEY (id),
+ UNIQUE KEY uq_camro_membership_resource (compute_allocation_membership_id, compute_allocation_resource_id),
+ KEY idx_camro_resource (compute_allocation_resource_id),
+ CONSTRAINT fk_camro_membership FOREIGN KEY (compute_allocation_membership_id)
+ REFERENCES compute_allocation_memberships (id) ON DELETE CASCADE,
+ CONSTRAINT fk_camro_resource FOREIGN KEY (compute_allocation_resource_id)
+ REFERENCES compute_allocation_resources (id) ON DELETE CASCADE
+) ENGINE = InnoDB DEFAULT CHARSET = utf8mb4 COLLATE = utf8mb4_unicode_ci;
diff --git a/internal/server/server.go b/internal/server/server.go
index 4b3e20fd2..960799fd8 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -111,11 +111,17 @@ func (s *Server) routes() {
s.mux.HandleFunc("POST /compute-allocation-memberships", s.createComputeAllocationMembership)
s.mux.HandleFunc("GET /compute-allocation-memberships/{id}", s.getComputeAllocationMembership)
s.mux.HandleFunc("PUT /compute-allocation-memberships/{id}", s.updateComputeAllocationMembership)
- s.mux.HandleFunc("PUT /compute-allocation-memberships/{id}/allocation-amount", s.updateMembershipAllocationAmount)
s.mux.HandleFunc("PUT /compute-allocation-memberships/{id}/status", s.updateMembershipStatus)
s.mux.HandleFunc("DELETE /compute-allocation-memberships/{id}", s.deleteComputeAllocationMembership)
s.mux.HandleFunc("GET /compute-allocations/{id}/memberships", s.listMembersForAllocation)
s.mux.HandleFunc("GET /users/{id}/compute-allocation-memberships", s.listAllocationsForUser)
+ s.mux.HandleFunc("GET /compute-allocation-memberships/{id}/resource-overrides", s.listOverridesForMembership)
+
+ s.mux.HandleFunc("POST /compute-allocation-membership-resource-overrides", s.createComputeAllocationMembershipResourceOverride)
+ s.mux.HandleFunc("GET /compute-allocation-membership-resource-overrides/{id}", s.getComputeAllocationMembershipResourceOverride)
+ s.mux.HandleFunc("PUT /compute-allocation-membership-resource-overrides/{id}", s.updateComputeAllocationMembershipResourceOverride)
+ s.mux.HandleFunc("DELETE /compute-allocation-membership-resource-overrides/{id}", s.deleteComputeAllocationMembershipResourceOverride)
+ s.mux.HandleFunc("GET /compute-allocation-resources/{id}/membership-overrides", s.listOverridesForResource)
s.mux.HandleFunc("POST /compute-allocation-usages", s.createComputeAllocationUsage)
s.mux.HandleFunc("GET /compute-allocation-usages/{id}", s.getComputeAllocationUsage)
@@ -652,22 +658,6 @@ func (s *Server) updateComputeAllocationMembership(w http.ResponseWriter, r *htt
writeJSON(w, http.StatusOK, updated)
}
-func (s *Server) updateMembershipAllocationAmount(w http.ResponseWriter, r *http.Request) {
- var body struct {
- AllocationAmount int64 `json:"allocation_amount"`
- }
- if err := decodeJSON(r, &body); err != nil {
- writeError(w, http.StatusBadRequest, err)
- return
- }
- updated, err := s.svc.UpdateMembershipAllocationAmount(r.Context(), r.PathValue("id"), body.AllocationAmount)
- if err != nil {
- writeServiceError(w, err)
- return
- }
- writeJSON(w, http.StatusOK, updated)
-}
-
func (s *Server) updateMembershipStatus(w http.ResponseWriter, r *http.Request) {
var body struct {
MembershipStatus models.AllocationStatus `json:"membership_status"`
@@ -759,6 +749,70 @@ func (s *Server) listUsagesByUser(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, rows)
}
+func (s *Server) createComputeAllocationMembershipResourceOverride(w http.ResponseWriter, r *http.Request) {
+ var o models.ComputeAllocationMembershipResourceOverride
+ if err := decodeJSON(r, &o); err != nil {
+ writeError(w, http.StatusBadRequest, err)
+ return
+ }
+ created, err := s.svc.CreateComputeAllocationMembershipResourceOverride(r.Context(), &o)
+ if err != nil {
+ writeServiceError(w, err)
+ return
+ }
+ writeJSON(w, http.StatusCreated, created)
+}
+
+func (s *Server) getComputeAllocationMembershipResourceOverride(w http.ResponseWriter, r *http.Request) {
+ o, err := s.svc.GetComputeAllocationMembershipResourceOverride(r.Context(), r.PathValue("id"))
+ if err != nil {
+ writeServiceError(w, err)
+ return
+ }
+ writeJSON(w, http.StatusOK, o)
+}
+
+func (s *Server) updateComputeAllocationMembershipResourceOverride(w http.ResponseWriter, r *http.Request) {
+ var o models.ComputeAllocationMembershipResourceOverride
+ if err := decodeJSON(r, &o); err != nil {
+ writeError(w, http.StatusBadRequest, err)
+ return
+ }
+ o.ID = r.PathValue("id")
+ updated, err := s.svc.UpdateComputeAllocationMembershipResourceOverride(r.Context(), &o)
+ if err != nil {
+ writeServiceError(w, err)
+ return
+ }
+ writeJSON(w, http.StatusOK, updated)
+}
+
+func (s *Server) deleteComputeAllocationMembershipResourceOverride(w http.ResponseWriter, r *http.Request) {
+ if err := s.svc.DeleteComputeAllocationMembershipResourceOverride(r.Context(), r.PathValue("id")); err != nil {
+ writeServiceError(w, err)
+ return
+ }
+ w.WriteHeader(http.StatusNoContent)
+}
+
+func (s *Server) listOverridesForMembership(w http.ResponseWriter, r *http.Request) {
+ rows, err := s.svc.ListOverridesForMembership(r.Context(), r.PathValue("id"))
+ if err != nil {
+ writeServiceError(w, err)
+ return
+ }
+ writeJSON(w, http.StatusOK, rows)
+}
+
+func (s *Server) listOverridesForResource(w http.ResponseWriter, r *http.Request) {
+ rows, err := s.svc.ListOverridesForResource(r.Context(), r.PathValue("id"))
+ if err != nil {
+ writeServiceError(w, err)
+ return
+ }
+ writeJSON(w, http.StatusOK, rows)
+}
+
func (s *Server) getTotalSUUsageForAllocation(w http.ResponseWriter, r *http.Request) {
total, err := s.svc.GetTotalSUUsageForAllocation(r.Context(), r.PathValue("id"))
if err != nil {
diff --git a/internal/store/compute_allocation_membership_resource_override_store.go b/internal/store/compute_allocation_membership_resource_override_store.go
new file mode 100644
index 000000000..1e0f0d1b0
--- /dev/null
+++ b/internal/store/compute_allocation_membership_resource_override_store.go
@@ -0,0 +1,123 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package store
+
+import (
+ "context"
+ "database/sql"
+ "errors"
+
+ "github.com/jmoiron/sqlx"
+
+ "github.com/apache/airavata-custos/pkg/models"
+)
+
+const computeAllocationMembershipResourceOverrideColumns = "id, compute_allocation_membership_id, compute_allocation_resource_id, overridden_resource_amount"
+
+type mysqlComputeAllocationMembershipResourceOverrideStore struct {
+ db *sqlx.DB
+}
+
+// NewComputeAllocationMembershipResourceOverrideStore returns a MySQL-backed
+// ComputeAllocationMembershipResourceOverrideStore.
+func NewComputeAllocationMembershipResourceOverrideStore(db *sqlx.DB) ComputeAllocationMembershipResourceOverrideStore {
+ return &mysqlComputeAllocationMembershipResourceOverrideStore{db: db}
+}
+
+func (s *mysqlComputeAllocationMembershipResourceOverrideStore) FindByID(ctx context.Context, id string) (*models.ComputeAllocationMembershipResourceOverride, error) {
+ var o models.ComputeAllocationMembershipResourceOverride
+ err := s.db.GetContext(ctx, &o,
+ `SELECT `+computeAllocationMembershipResourceOverrideColumns+`
+ FROM compute_allocation_membership_resource_overrides
+ WHERE id = ?`, id)
+ if err != nil {
+ if errors.Is(err, sql.ErrNoRows) {
+ return nil, nil
+ }
+ return nil, err
+ }
+ return &o, nil
+}
+
+func (s *mysqlComputeAllocationMembershipResourceOverrideStore) FindByPair(ctx context.Context, membershipID, resourceID string) (*models.ComputeAllocationMembershipResourceOverride, error) {
+ var o models.ComputeAllocationMembershipResourceOverride
+ err := s.db.GetContext(ctx, &o,
+ `SELECT `+computeAllocationMembershipResourceOverrideColumns+`
+ FROM compute_allocation_membership_resource_overrides
+ WHERE compute_allocation_membership_id = ? AND compute_allocation_resource_id = ?`,
+ membershipID, resourceID)
+ if err != nil {
+ if errors.Is(err, sql.ErrNoRows) {
+ return nil, nil
+ }
+ return nil, err
+ }
+ return &o, nil
+}
+
+func (s *mysqlComputeAllocationMembershipResourceOverrideStore) FindByMembership(ctx context.Context, membershipID string) ([]models.ComputeAllocationMembershipResourceOverride, error) {
+ var rows []models.ComputeAllocationMembershipResourceOverride
+ err := s.db.SelectContext(ctx, &rows,
+ `SELECT `+computeAllocationMembershipResourceOverrideColumns+`
+ FROM compute_allocation_membership_resource_overrides
+ WHERE compute_allocation_membership_id = ?
+ ORDER BY compute_allocation_resource_id`, membershipID)
+ if err != nil {
+ return nil, err
+ }
+ return rows, nil
+}
+
+func (s *mysqlComputeAllocationMembershipResourceOverrideStore) FindByResource(ctx context.Context, resourceID string) ([]models.ComputeAllocationMembershipResourceOverride, error) {
+ var rows []models.ComputeAllocationMembershipResourceOverride
+ err := s.db.SelectContext(ctx, &rows,
+ `SELECT `+computeAllocationMembershipResourceOverrideColumns+`
+ FROM compute_allocation_membership_resource_overrides
+ WHERE compute_allocation_resource_id = ?
+ ORDER BY compute_allocation_membership_id`, resourceID)
+ if err != nil {
+ return nil, err
+ }
+ return rows, nil
+}
+
+func (s *mysqlComputeAllocationMembershipResourceOverrideStore) Create(ctx context.Context, tx *sql.Tx, o *models.ComputeAllocationMembershipResourceOverride) error {
+ _, err := tx.ExecContext(ctx,
+ `INSERT INTO compute_allocation_membership_resource_overrides
+ (id, compute_allocation_membership_id, compute_allocation_resource_id, overridden_resource_amount)
+ VALUES (?, ?, ?, ?)`,
+ o.ID, o.ComputeAllocationMembershipID, o.ComputeAllocationResourceID, o.OverriddenResourceAmount)
+ return err
+}
+
+func (s *mysqlComputeAllocationMembershipResourceOverrideStore) Update(ctx context.Context, tx *sql.Tx, o *models.ComputeAllocationMembershipResourceOverride) error {
+ _, err := tx.ExecContext(ctx,
+ `UPDATE compute_allocation_membership_resource_overrides
+ SET compute_allocation_membership_id = ?,
+ compute_allocation_resource_id = ?,
+ overridden_resource_amount = ?
+ WHERE id = ?`,
+ o.ComputeAllocationMembershipID, o.ComputeAllocationResourceID, o.OverriddenResourceAmount, o.ID)
+ return err
+}
+
+func (s *mysqlComputeAllocationMembershipResourceOverrideStore) Delete(ctx context.Context, tx *sql.Tx, id string) error {
+ _, err := tx.ExecContext(ctx,
+ `DELETE FROM compute_allocation_membership_resource_overrides WHERE id = ?`, id)
+ return err
+}
diff --git a/internal/store/compute_allocation_membership_store.go b/internal/store/compute_allocation_membership_store.go
index aae350ac9..baf217782 100644
--- a/internal/store/compute_allocation_membership_store.go
+++ b/internal/store/compute_allocation_membership_store.go
@@ -27,7 +27,7 @@ import (
"github.com/apache/airavata-custos/pkg/models"
)
-const computeAllocationMembershipColumns = "id, compute_allocation_id, user_id, allocation_amount, start_time, end_time, membership_status"
+const computeAllocationMembershipColumns = "id, compute_allocation_id, user_id, start_time, end_time, membership_status"
type mysqlComputeAllocationMembershipStore struct {
db *sqlx.DB
@@ -96,9 +96,9 @@ func (s *mysqlComputeAllocationMembershipStore) FindByUser(ctx context.Context,
func (s *mysqlComputeAllocationMembershipStore) Create(ctx context.Context, tx *sql.Tx, m *models.ComputeAllocationMembership) error {
_, err := tx.ExecContext(ctx,
`INSERT INTO compute_allocation_memberships
- (id, compute_allocation_id, user_id, allocation_amount, start_time, end_time, membership_status)
- VALUES (?, ?, ?, ?, ?, ?, ?)`,
- m.ID, m.ComputeAllocationID, m.UserID, m.AllocationAmount, m.StartTime, m.EndTime, string(m.MembershipStatus))
+ (id, compute_allocation_id, user_id, start_time, end_time, membership_status)
+ VALUES (?, ?, ?, ?, ?, ?)`,
+ m.ID, m.ComputeAllocationID, m.UserID, m.StartTime, m.EndTime, string(m.MembershipStatus))
return err
}
@@ -107,12 +107,11 @@ func (s *mysqlComputeAllocationMembershipStore) Update(ctx context.Context, tx *
`UPDATE compute_allocation_memberships
SET compute_allocation_id = ?,
user_id = ?,
- allocation_amount = ?,
start_time = ?,
end_time = ?,
membership_status = ?
WHERE id = ?`,
- m.ComputeAllocationID, m.UserID, m.AllocationAmount, m.StartTime, m.EndTime, string(m.MembershipStatus), m.ID)
+ m.ComputeAllocationID, m.UserID, m.StartTime, m.EndTime, string(m.MembershipStatus), m.ID)
return err
}
diff --git a/internal/store/store.go b/internal/store/store.go
index 4dddec180..8066d87e8 100644
--- a/internal/store/store.go
+++ b/internal/store/store.go
@@ -250,6 +250,26 @@ type ComputeAllocationMembershipStore interface {
Delete(ctx context.Context, tx *sql.Tx, id string) error
}
+// ComputeAllocationMembershipResourceOverrideStore defines persistence
+// operations for per-resource overrides of the SU amount granted to a
+// compute allocation membership.
+type ComputeAllocationMembershipResourceOverrideStore interface {
+ // FindByID returns the override with the given ID, or nil if it does not exist.
+ FindByID(ctx context.Context, id string) (*models.ComputeAllocationMembershipResourceOverride, error)
+ // FindByPair returns the override for a (membership, resource) pair, or nil if absent.
+ FindByPair(ctx context.Context, membershipID, resourceID string) (*models.ComputeAllocationMembershipResourceOverride, error)
+ // FindByMembership returns every override recorded against the given membership.
+ FindByMembership(ctx context.Context, membershipID string) ([]models.ComputeAllocationMembershipResourceOverride, error)
+ // FindByResource returns every override referencing the given resource.
+ FindByResource(ctx context.Context, resourceID string) ([]models.ComputeAllocationMembershipResourceOverride, error)
+ // Create inserts a new override within the provided transaction.
+ Create(ctx context.Context, tx *sql.Tx, o *models.ComputeAllocationMembershipResourceOverride) error
+ // Update replaces mutable fields of an existing override within the provided transaction.
+ Update(ctx context.Context, tx *sql.Tx, o *models.ComputeAllocationMembershipResourceOverride) error
+ // Delete removes an override by ID within the provided transaction.
+ Delete(ctx context.Context, tx *sql.Tx, id string) error
+}
+
// ComputeAllocationUsageStore defines persistence operations for the
// append-only log of resource consumption events charged against a compute
// allocation.
diff --git a/pkg/events/compute_allocation_membership_resource_override_subscribe.go b/pkg/events/compute_allocation_membership_resource_override_subscribe.go
new file mode 100644
index 000000000..ea77026fa
--- /dev/null
+++ b/pkg/events/compute_allocation_membership_resource_override_subscribe.go
@@ -0,0 +1,67 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package events
+
+import (
+ "log/slog"
+
+ "github.com/apache/airavata-custos/pkg/models"
+)
+
+// ComputeAllocationMembershipResourceOverrideHandler handles lifecycle
+// events for membership resource overrides with a typed payload.
+type ComputeAllocationMembershipResourceOverrideHandler func(o models.ComputeAllocationMembershipResourceOverride)
+
+// SubscribeComputeAllocationMembershipResourceOverrideCreated registers a
+// typed handler invoked whenever a
+// compute_allocation_membership_resource_override::create event is published.
+func (b *Bus) SubscribeComputeAllocationMembershipResourceOverrideCreated(handler ComputeAllocationMembershipResourceOverrideHandler) {
+ b.subscribeMembershipResourceOverride(ComputeAllocationMembershipResourceOverrideCreateEvent, handler)
+}
+
+// SubscribeComputeAllocationMembershipResourceOverrideUpdated registers a
+// typed handler invoked whenever a
+// compute_allocation_membership_resource_override::update event is published.
+func (b *Bus) SubscribeComputeAllocationMembershipResourceOverrideUpdated(handler ComputeAllocationMembershipResourceOverrideHandler) {
+ b.subscribeMembershipResourceOverride(ComputeAllocationMembershipResourceOverrideUpdateEvent, handler)
+}
+
+// SubscribeComputeAllocationMembershipResourceOverrideDeleted registers a
+// typed handler invoked whenever a
+// compute_allocation_membership_resource_override::delete event is published.
+func (b *Bus) SubscribeComputeAllocationMembershipResourceOverrideDeleted(handler ComputeAllocationMembershipResourceOverrideHandler) {
+ b.subscribeMembershipResourceOverride(ComputeAllocationMembershipResourceOverrideDeleteEvent, handler)
+}
+
+func (b *Bus) subscribeMembershipResourceOverride(topic EventType, handler ComputeAllocationMembershipResourceOverrideHandler) {
+ b.Subscribe(topic, func(event Event, value interface{}) {
+ switch o := value.(type) {
+ case models.ComputeAllocationMembershipResourceOverride:
+ handler(o)
+ case *models.ComputeAllocationMembershipResourceOverride:
+ if o != nil {
+ handler(*o)
+ }
+ default:
+ slog.Warn("compute allocation membership resource override event payload has unexpected type",
+ "type", event.Type,
+ "got", value,
+ )
+ }
+ })
+}
diff --git a/pkg/events/types.go b/pkg/events/types.go
index 44758f0e4..8ae8de802 100644
--- a/pkg/events/types.go
+++ b/pkg/events/types.go
@@ -96,6 +96,13 @@ const (
ComputeAllocationMembershipDeleteEvent EventType = "compute_allocation_membership::delete"
)
+// ComputeAllocationMembershipResourceOverride lifecycle message types.
+const (
+ ComputeAllocationMembershipResourceOverrideCreateEvent EventType = "compute_allocation_membership_resource_override::create"
+ ComputeAllocationMembershipResourceOverrideUpdateEvent EventType = "compute_allocation_membership_resource_override::update"
+ ComputeAllocationMembershipResourceOverrideDeleteEvent EventType = "compute_allocation_membership_resource_override::delete"
+)
+
// ComputeAllocationResourceMapping lifecycle message types.
const (
ComputeAllocationResourceMappingCreateEvent EventType = "compute_allocation_resource_mapping::create"
diff --git a/pkg/models/allocation.go b/pkg/models/allocation.go
index 7813dcc2a..ca95595b5 100644
--- a/pkg/models/allocation.go
+++ b/pkg/models/allocation.go
@@ -35,9 +35,9 @@ type ComputeAllocation struct {
type ComputeAllocationResource struct {
ID string `json:"id" db:"id"`
- Name string `json:"name" db:"name"` // A human-readable name for the resource, e.g., "GPU B200", "CPU", "GPU RTX6000", etc.
- ResourceType string `json:"resource_type" db:"resource_type"` // CPU, GPU, etc.
- ResourceAmount int64 `json:"resource_amount" db:"resource_amount"` // Number of CPUs, GPUs, etc. allocated.
+ Name string `json:"name" db:"name"` // resource / partition name, e.g., "cpu", "gpu", etc.
+ ResourceType string `json:"resource_type" db:"resource_type"` // GrpTRES, GrpTRESMins
+ ResourceAmount int64 `json:"resource_amount" db:"resource_amount"` // Number of CPUs, GPUs, time in minutes, or other unit depending on the resource type.
}
type ComputeAllocationResourceMapping struct {
@@ -95,11 +95,17 @@ type ComputeAllocationUsage struct { // Represents the usage of a compute alloca
ComputeAllocationResourceID string `json:"compute_allocation_resource_id" db:"compute_allocation_resource_id"` // The specific resource consumed, e.g., 20 CPU hours, 10 GPU hours, etc.
}
+type ComputeAllocationMembershipResourceOverride struct {
+ ID string `json:"id" db:"id"`
+ ComputeAllocationMembershipID string `json:"compute_allocation_membership_id" db:"compute_allocation_membership_id"`
+ ComputeAllocationResourceID string `json:"compute_allocation_resource_id" db:"compute_allocation_resource_id"`
+ OverriddenResourceAmount int64 `json:"overridden_resource_amount" db:"overridden_resource_amount"` // The overridden SU amount for the user for this specific resource, e.g., 150 SUs for GPU hours, etc.
+}
+
type ComputeAllocationMembership struct {
ID string `json:"id" db:"id"`
ComputeAllocationID string `json:"compute_allocation_id" db:"compute_allocation_id"`
UserID string `json:"user_id" db:"user_id"`
- AllocationAmount int64 `json:"allocation_amount" db:"allocation_amount"` // SUs allocated to the user, e.g., 100 CPU hours, 50 GPU hours, etc.
StartTime time.Time `json:"start_time" db:"start_time"`
EndTime time.Time `json:"end_time" db:"end_time"`
MembershipStatus AllocationStatus `json:"membership_status" db:"membership_status"` // ACTIVE, INACTIVE, etc.
diff --git a/pkg/service/compute_allocation_membership.go b/pkg/service/compute_allocation_membership.go
index 9d6e0056d..f61bbdc06 100644
--- a/pkg/service/compute_allocation_membership.go
+++ b/pkg/service/compute_allocation_membership.go
@@ -154,33 +154,6 @@ func (s *Service) UpdateComputeAllocationMembership(ctx context.Context, m *mode
return m, nil
}
-// UpdateMembershipAllocationAmount sets the SU sub-allocation granted to the
-// user identified by the given membership ID. Other fields are preserved.
-func (s *Service) UpdateMembershipAllocationAmount(ctx context.Context, id string, amount int64) (*models.ComputeAllocationMembership, error) {
- if id == "" {
- return nil, fmt.Errorf("%w: compute allocation membership id is required", ErrInvalidInput)
- }
- if amount < 0 {
- return nil, fmt.Errorf("%w: allocation_amount must be non-negative", ErrInvalidInput)
- }
- existing, err := s.memberships.FindByID(ctx, id)
- if err != nil {
- return nil, fmt.Errorf("lookup compute allocation membership: %w", err)
- }
- if existing == nil {
- return nil, ErrNotFound
- }
- existing.AllocationAmount = amount
- if err := s.inTx(ctx, func(tx *sql.Tx) error {
- return s.memberships.Update(ctx, tx, existing)
- }); err != nil {
- return nil, fmt.Errorf("update compute allocation membership amount: %w", err)
- }
-
- s.eventBus.Publish(events.ComputeAllocationMembershipUpdateEvent, existing)
- return existing, nil
-}
-
// UpdateMembershipStatus sets the lifecycle status (ACTIVE, INACTIVE, etc.)
// of the membership identified by the given ID. Other fields are preserved.
func (s *Service) UpdateMembershipStatus(ctx context.Context, id string, status models.AllocationStatus) (*models.ComputeAllocationMembership, error) {
diff --git a/pkg/service/compute_allocation_membership_resource_override.go b/pkg/service/compute_allocation_membership_resource_override.go
new file mode 100644
index 000000000..81fadee40
--- /dev/null
+++ b/pkg/service/compute_allocation_membership_resource_override.go
@@ -0,0 +1,191 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package service
+
+import (
+ "context"
+ "database/sql"
+ "fmt"
+
+ "github.com/apache/airavata-custos/pkg/events"
+ "github.com/apache/airavata-custos/pkg/models"
+)
+
+// CreateComputeAllocationMembershipResourceOverride records a per-resource
+// override of the SU amount granted to the membership. The referenced
+// membership and resource must exist, and the (membership, resource) pair
+// must not already have an override.
+func (s *Service) CreateComputeAllocationMembershipResourceOverride(ctx context.Context, o *models.ComputeAllocationMembershipResourceOverride) (*models.ComputeAllocationMembershipResourceOverride, error) {
+ if o == nil {
+ return nil, fmt.Errorf("%w: membership resource override is nil", ErrInvalidInput)
+ }
+ if o.ComputeAllocationMembershipID == "" {
+ return nil, fmt.Errorf("%w: compute_allocation_membership_id is required", ErrInvalidInput)
+ }
+ if o.ComputeAllocationResourceID == "" {
+ return nil, fmt.Errorf("%w: compute_allocation_resource_id is required", ErrInvalidInput)
+ }
+ if o.OverriddenResourceAmount < 0 {
+ return nil, fmt.Errorf("%w: overridden_resource_amount must be non-negative", ErrInvalidInput)
+ }
+
+ if m, err := s.memberships.FindByID(ctx, o.ComputeAllocationMembershipID); err != nil {
+ return nil, fmt.Errorf("lookup compute allocation membership: %w", err)
+ } else if m == nil {
+ return nil, fmt.Errorf("%w: compute allocation membership %q not found",
+ ErrInvalidInput, o.ComputeAllocationMembershipID)
+ }
+ if r, err := s.resources.FindByID(ctx, o.ComputeAllocationResourceID); err != nil {
+ return nil, fmt.Errorf("lookup compute allocation resource: %w", err)
+ } else if r == nil {
+ return nil, fmt.Errorf("%w: compute allocation resource %q not found",
+ ErrInvalidInput, o.ComputeAllocationResourceID)
+ }
+
+ if existing, err := s.membershipOverrides.FindByPair(ctx, o.ComputeAllocationMembershipID, o.ComputeAllocationResourceID); err != nil {
+ return nil, fmt.Errorf("lookup membership resource override: %w", err)
+ } else if existing != nil {
+ return nil, fmt.Errorf("%w: override already exists for membership %q resource %q",
+ ErrAlreadyExists, o.ComputeAllocationMembershipID, o.ComputeAllocationResourceID)
+ }
+
+ if o.ID == "" {
+ o.ID = newID()
+ }
+
+ if err := s.inTx(ctx, func(tx *sql.Tx) error {
+ return s.membershipOverrides.Create(ctx, tx, o)
+ }); err != nil {
+ return nil, fmt.Errorf("create membership resource override: %w", err)
+ }
+
+ s.eventBus.Publish(events.ComputeAllocationMembershipResourceOverrideCreateEvent, o)
+ return o, nil
+}
+
+// GetComputeAllocationMembershipResourceOverride retrieves an override by ID.
+func (s *Service) GetComputeAllocationMembershipResourceOverride(ctx context.Context, id string) (*models.ComputeAllocationMembershipResourceOverride, error) {
+ o, err := s.membershipOverrides.FindByID(ctx, id)
+ if err != nil {
+ return nil, fmt.Errorf("get membership resource override: %w", err)
+ }
+ if o == nil {
+ return nil, ErrNotFound
+ }
+ return o, nil
+}
+
+// GetComputeAllocationMembershipResourceOverrideByPair retrieves the override
+// for the given (membership, resource) pair.
+func (s *Service) GetComputeAllocationMembershipResourceOverrideByPair(ctx context.Context, membershipID, resourceID string) (*models.ComputeAllocationMembershipResourceOverride, error) {
+ if membershipID == "" {
+ return nil, fmt.Errorf("%w: compute_allocation_membership_id is required", ErrInvalidInput)
+ }
+ if resourceID == "" {
+ return nil, fmt.Errorf("%w: compute_allocation_resource_id is required", ErrInvalidInput)
+ }
+ o, err := s.membershipOverrides.FindByPair(ctx, membershipID, resourceID)
+ if err != nil {
+ return nil, fmt.Errorf("get membership resource override by pair: %w", err)
+ }
+ if o == nil {
+ return nil, ErrNotFound
+ }
+ return o, nil
+}
+
+// ListOverridesForMembership returns every resource override recorded
+// against the given membership.
+func (s *Service) ListOverridesForMembership(ctx context.Context, membershipID string) ([]models.ComputeAllocationMembershipResourceOverride, error) {
+ if membershipID == "" {
+ return nil, fmt.Errorf("%w: compute_allocation_membership_id is required", ErrInvalidInput)
+ }
+ rows, err := s.membershipOverrides.FindByMembership(ctx, membershipID)
+ if err != nil {
+ return nil, fmt.Errorf("list overrides for membership: %w", err)
+ }
+ return rows, nil
+}
+
+// ListOverridesForResource returns every membership override referencing the
+// given resource.
+func (s *Service) ListOverridesForResource(ctx context.Context, resourceID string) ([]models.ComputeAllocationMembershipResourceOverride, error) {
+ if resourceID == "" {
+ return nil, fmt.Errorf("%w: compute_allocation_resource_id is required", ErrInvalidInput)
+ }
+ rows, err := s.membershipOverrides.FindByResource(ctx, resourceID)
+ if err != nil {
+ return nil, fmt.Errorf("list overrides for resource: %w", err)
+ }
+ return rows, nil
+}
+
+// UpdateComputeAllocationMembershipResourceOverride replaces mutable fields
+// of an existing override.
+func (s *Service) UpdateComputeAllocationMembershipResourceOverride(ctx context.Context, o *models.ComputeAllocationMembershipResourceOverride) (*models.ComputeAllocationMembershipResourceOverride, error) {
+ if o == nil || o.ID == "" {
+ return nil, fmt.Errorf("%w: membership resource override id is required", ErrInvalidInput)
+ }
+ if o.OverriddenResourceAmount < 0 {
+ return nil, fmt.Errorf("%w: overridden_resource_amount must be non-negative", ErrInvalidInput)
+ }
+ existing, err := s.membershipOverrides.FindByID(ctx, o.ID)
+ if err != nil {
+ return nil, fmt.Errorf("lookup membership resource override: %w", err)
+ }
+ if existing == nil {
+ return nil, ErrNotFound
+ }
+ if o.ComputeAllocationMembershipID == "" {
+ o.ComputeAllocationMembershipID = existing.ComputeAllocationMembershipID
+ }
+ if o.ComputeAllocationResourceID == "" {
+ o.ComputeAllocationResourceID = existing.ComputeAllocationResourceID
+ }
+
+ if err := s.inTx(ctx, func(tx *sql.Tx) error {
+ return s.membershipOverrides.Update(ctx, tx, o)
+ }); err != nil {
+ return nil, fmt.Errorf("update membership resource override: %w", err)
+ }
+
+ s.eventBus.Publish(events.ComputeAllocationMembershipResourceOverrideUpdateEvent, o)
+ return o, nil
+}
+
+// DeleteComputeAllocationMembershipResourceOverride removes an override by ID.
+func (s *Service) DeleteComputeAllocationMembershipResourceOverride(ctx context.Context, id string) error {
+ if id == "" {
+ return fmt.Errorf("%w: membership resource override id is required", ErrInvalidInput)
+ }
+ existing, err := s.membershipOverrides.FindByID(ctx, id)
+ if err != nil {
+ return fmt.Errorf("lookup membership resource override: %w", err)
+ }
+ if existing == nil {
+ return ErrNotFound
+ }
+ if err := s.inTx(ctx, func(tx *sql.Tx) error {
+ return s.membershipOverrides.Delete(ctx, tx, id)
+ }); err != nil {
+ return fmt.Errorf("delete membership resource override: %w", err)
+ }
+
+ s.eventBus.Publish(events.ComputeAllocationMembershipResourceOverrideDeleteEvent, existing)
+ return nil
+}
diff --git a/pkg/service/service.go b/pkg/service/service.go
index e187ff82f..23ca9bc83 100644
--- a/pkg/service/service.go
+++ b/pkg/service/service.go
@@ -49,6 +49,7 @@ type Service struct {
changeRequests store.ComputeAllocationChangeRequestStore
changeEvents store.ComputeAllocationChangeRequestEventStore
memberships store.ComputeAllocationMembershipStore
+ membershipOverrides store.ComputeAllocationMembershipResourceOverrideStore
usages store.ComputeAllocationUsageStore
}
@@ -71,6 +72,7 @@ func New(database *sqlx.DB, eventBus *events.Bus) *Service {
changeRequests: store.NewComputeAllocationChangeRequestStore(database),
changeEvents: store.NewComputeAllocationChangeRequestEventStore(database),
memberships: store.NewComputeAllocationMembershipStore(database),
+ membershipOverrides: store.NewComputeAllocationMembershipResourceOverrideStore(database),
usages: store.NewComputeAllocationUsageStore(database),
}
}
@@ -93,6 +95,7 @@ func NewWithStores(
allocDiffs store.ComputeAllocationDiffStore,
changeRequests store.ComputeAllocationChangeRequestStore,
changeEvents store.ComputeAllocationChangeRequestEventStore,
+ membershipOverrides store.ComputeAllocationMembershipResourceOverrideStore,
memberships store.ComputeAllocationMembershipStore,
usages store.ComputeAllocationUsageStore,
) *Service {
@@ -111,6 +114,7 @@ func NewWithStores(
allocDiffs: allocDiffs,
changeRequests: changeRequests,
changeEvents: changeEvents,
+ membershipOverrides: membershipOverrides,
memberships: memberships,
usages: usages,
}
(airavata-custos) 04/07: Adding SLURM association when an allocation membership is created
Posted by di...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
DImuthuUpe pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/airavata-custos.git
commit 6b7e6daa6e18ca5da1233ae5bfc4af34f059eeb2
Author: DImuthuUpe <di...@gmail.com>
AuthorDate: Mon May 18 23:47:59 2026 -0400
Adding SLURM association when an allocation membership is created
---
.../internal/subscribers/members.go | 64 +++++++
.../internal/subscribers/subscriber.go | 1 +
docs/API-Docs.md | 195 ++++++++++++++++++++-
.../000010_compute_cluster_users.down.sql | 18 ++
.../migrations/000010_compute_cluster_users.up.sql | 31 ++++
internal/server/server.go | 80 +++++++++
internal/store/compute_cluster_user_store.go | 116 ++++++++++++
internal/store/store.go | 19 ++
pkg/events/compute_cluster_user_subscribe.go | 64 +++++++
pkg/events/types.go | 7 +
pkg/models/allocation.go | 7 +
pkg/service/compute_cluster_user.go | 179 +++++++++++++++++++
pkg/service/service.go | 4 +
13 files changed, 784 insertions(+), 1 deletion(-)
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/members.go b/connectors/SLURM/Association-Mapper/internal/subscribers/members.go
new file mode 100644
index 000000000..cd609d5af
--- /dev/null
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/members.go
@@ -0,0 +1,64 @@
+package subscribers
+
+import (
+ "log/slog"
+
+ "context"
+ client "github.com/apache/airavata-custos/connectors/SLURM/Association-Mapper/internal/operations"
+ "github.com/apache/airavata-custos/pkg/models"
+ "time"
+)
+
+func (a *AssociationSubscriber) SubscribeToComputeAllocationMembershipCreation(membership models.ComputeAllocationMembership) {
+ slog.Info("Received compute allocation membership creation event", "membership", membership)
+
+ ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+ defer cancel()
+
+ allocation, err := a.coreService.GetComputeAllocation(ctx, membership.ComputeAllocationID)
+ if err != nil {
+ slog.Error("Failed to get compute allocation", "error", err)
+ return
+ }
+
+ cluster, err := a.coreService.GetComputeCluster(ctx, allocation.ComputeClusterID)
+ if err != nil {
+ slog.Error("Failed to get compute cluster", "error", err)
+ return
+ }
+
+ user, err := a.coreService.GetUser(ctx, membership.UserID)
+ if err != nil {
+ slog.Error("Failed to get user", "error", err)
+ return
+ }
+
+ csu, err := a.coreService.GetComputeClusterUserByPair(ctx, cluster.ID, user.ID) // TODO: use this to get the local username for the association instead of assuming it's the same as the Airavata Custos username
+ if err != nil {
+ slog.Error("Failed to get compute cluster user by pair", "error", err)
+ return
+ }
+
+ resources, err := a.coreService.ListResourcesForAllocation(ctx, allocation.ID) // TODO: use this to get the partition for the association instead of hardcoding it to "default"
+ if err != nil {
+ slog.Error("Failed to list resources for allocation", "error", err)
+ return
+ }
+
+ if len(resources) == 0 {
+ slog.Warn("No resources found for allocation, defaulting to partition 'default'", "allocation_id", allocation.ID)
+ }
+ association := client.Association{
+ Account: allocation.Name,
+ Cluster: cluster.Name,
+ User: csu.LocalUsername,
+ Partition: resources[0].Name,
+ }
+
+ err = a.slurmClient.UpsertAssociation(association)
+ if err != nil {
+ slog.Error("Failed to upsert association", "error", err)
+ } else {
+ slog.Info("Successfully upserted association", "association", association)
+ }
+}
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go b/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go
index ce5161fd7..658dcc073 100644
--- a/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go
@@ -22,4 +22,5 @@ func (a *AssociationSubscriber) RegisterSubscribers() {
a.eventBus.SubscribeComputeAllocationCreated(a.SubscribeToComputeAllocationCreation)
a.eventBus.SubscribeComputeAllocationDeleted(a.SubscribeToComputeAllocationDeletion)
a.eventBus.SubscribeComputeAllocationUpdated(a.SubscribeToComputeAllocationUpdate)
+ a.eventBus.SubscribeComputeAllocationMembershipCreated(a.SubscribeToComputeAllocationMembershipCreation)
}
diff --git a/docs/API-Docs.md b/docs/API-Docs.md
index 6931d0304..81aa716db 100644
--- a/docs/API-Docs.md
+++ b/docs/API-Docs.md
@@ -330,6 +330,159 @@ Retrieve a single compute cluster by its ID.
---
+## Compute Cluster Users
+
+A compute-cluster user maps a Custos user to their local account
+(`local_username`) on a specific compute cluster. Each `(compute_cluster_id,
+user_id)` pair is unique. The mapping is removed automatically when either
+the referenced compute cluster or user is deleted.
+
+### `POST /compute-cluster-users`
+
+Create a new compute-cluster user mapping.
+
+**Required fields:** `compute_cluster_id`, `user_id`, `local_username`
+**Optional fields:** `id` (auto-generated if omitted)
+
+Both `compute_cluster_id` and `user_id` must reference existing records.
+
+**Request**
+
+```json
+{
+ "compute_cluster_id": "9b0a7f1c-2c5d-4e1b-9a0f-22e8a5c2dcb1",
+ "user_id": "5e2c7b3a-1d8f-4d2c-bf09-83a7c4d6e210",
+ "local_username": "jdoe"
+}
+```
+
+**Response 201**
+
+```json
+{
+ "id": "0d72d3b1-6f1a-4a92-9c4e-1d7a4b5f9c2d",
+ "compute_cluster_id": "9b0a7f1c-2c5d-4e1b-9a0f-22e8a5c2dcb1",
+ "user_id": "5e2c7b3a-1d8f-4d2c-bf09-83a7c4d6e210",
+ "local_username": "jdoe"
+}
+```
+
+**Errors**
+
+- `400` — `compute_cluster_id`, `user_id`, or `local_username` is missing,
+ or the referenced cluster/user does not exist.
+- `409` — this user is already mapped on the given compute cluster.
+
+---
+
+### `GET /compute-cluster-users/{id}`
+
+Retrieve a single compute-cluster user mapping by its ID.
+
+**Errors**
+
+- `404` — no mapping matches the supplied ID.
+
+---
+
+### `PUT /compute-cluster-users/{id}`
+
+Replace mutable fields of an existing compute-cluster user mapping. The
+path `{id}` overrides any `id` in the request body.
+
+**Required fields:** `compute_cluster_id`, `user_id`, `local_username`
+
+**Request**
+
+```json
+{
+ "compute_cluster_id": "9b0a7f1c-2c5d-4e1b-9a0f-22e8a5c2dcb1",
+ "user_id": "5e2c7b3a-1d8f-4d2c-bf09-83a7c4d6e210",
+ "local_username": "jane.doe"
+}
+```
+
+**Errors**
+
+- `400` — required fields missing.
+
+---
+
+### `DELETE /compute-cluster-users/{id}`
+
+Remove a compute-cluster user mapping.
+
+**Response 204** — empty body.
+
+**Errors**
+
+- `404` — no mapping matches the supplied ID.
+
+---
+
+### `GET /compute-clusters/{id}/users`
+
+List every user mapping for the given compute cluster, ordered by
+`local_username`.
+
+**Response 200**
+
+```json
+[
+ {
+ "id": "0d72d3b1-6f1a-4a92-9c4e-1d7a4b5f9c2d",
+ "compute_cluster_id": "9b0a7f1c-2c5d-4e1b-9a0f-22e8a5c2dcb1",
+ "user_id": "5e2c7b3a-1d8f-4d2c-bf09-83a7c4d6e210",
+ "local_username": "jdoe"
+ }
+]
+```
+
+---
+
+### `GET /compute-clusters/{id}/users/{userId}`
+
+Look up the single compute-cluster user mapping for the given
+`(compute_cluster_id, user_id)` pair.
+
+**Response 200**
+
+```json
+{
+ "id": "0d72d3b1-6f1a-4a92-9c4e-1d7a4b5f9c2d",
+ "compute_cluster_id": "9b0a7f1c-2c5d-4e1b-9a0f-22e8a5c2dcb1",
+ "user_id": "5e2c7b3a-1d8f-4d2c-bf09-83a7c4d6e210",
+ "local_username": "jdoe"
+}
+```
+
+**Errors**
+
+- `400` — `compute_cluster_id` or `user_id` is missing.
+- `404` — no mapping exists for the given pair.
+
+---
+
+### `GET /users/{id}/compute-cluster-users`
+
+List every cluster mapping held by the given Custos user, ordered by
+`compute_cluster_id`.
+
+**Response 200**
+
+```json
+[
+ {
+ "id": "0d72d3b1-6f1a-4a92-9c4e-1d7a4b5f9c2d",
+ "compute_cluster_id": "9b0a7f1c-2c5d-4e1b-9a0f-22e8a5c2dcb1",
+ "user_id": "5e2c7b3a-1d8f-4d2c-bf09-83a7c4d6e210",
+ "local_username": "jdoe"
+ }
+]
+```
+
+---
+
## Compute Allocations
A compute allocation grants a project a budget of Service Units (SUs) on a
@@ -1076,7 +1229,7 @@ ALLOC_ID=$(curl -s -X POST $BASE/compute-allocations \
RES_ID=$(curl -s -X POST $BASE/compute-allocation-resources \
-H 'Content-Type: application/json' \
- -d '{"name":"GPU B200","resource_type":"GPU","resource_amount":8}' | jq -r .id)
+ -d '{"name":"debug","resource_type":"CPU","resource_amount":24}' | jq -r .id)
# Attach the resource to the allocation.
curl -s -X POST $BASE/compute-allocations/$ALLOC_ID/resources \
@@ -1096,6 +1249,46 @@ curl -s -X POST $BASE/compute-allocation-resource-rates \
# Look up the currently-effective rate.
curl -s $BASE/compute-allocation-resources/$RES_ID/rates/effective | jq
+# Create a Custos user that will be granted access on the cluster.
+CLUSTER_USER_ACCT_ID=$(curl -s -X POST $BASE/users \
+ -H 'Content-Type: application/json' \
+ -d "{\"organization_id\":\"$ORG_ID\",\"first_name\":\"Dimuthu\",\"last_name\":\"Wannipurage\",\"email\":\"dimuthu@example.edu\"}" \
+ | jq -r .id)
+
+# Map the Custos user to a local UNIX account on the cluster.
+CLUSTER_USER_ID=$(curl -s -X POST $BASE/compute-cluster-users \
+ -H 'Content-Type: application/json' \
+ -d "{\"compute_cluster_id\":\"$CLUSTER_ID\",\"user_id\":\"$CLUSTER_USER_ACCT_ID\",\"local_username\":\"dimuthu\"}" \
+ | jq -r .id)
+
+# List all user mappings on the cluster.
+curl -s $BASE/compute-clusters/$CLUSTER_ID/users | jq
+
+# Grant the user a sub-allocation (membership) on the compute allocation.
+MEMBERSHIP_ID=$(curl -s -X POST $BASE/compute-allocation-memberships \
+ -H 'Content-Type: application/json' \
+ -d "{
+ \"compute_allocation_id\":\"$ALLOC_ID\",
+ \"user_id\":\"$CLUSTER_USER_ACCT_ID\",
+ \"allocation_amount\":50000,
+ \"start_time\":\"2026-04-01T00:00:00Z\",
+ \"end_time\":\"2026-06-30T23:59:59Z\",
+ \"membership_status\":\"ACTIVE\"
+ }" | jq -r .id)
+
+# List every member of the allocation.
+curl -s $BASE/compute-allocations/$ALLOC_ID/memberships | jq
+
+# List every allocation this user is a member of.
+curl -s $BASE/users/$CLUSTER_USER_ACCT_ID/compute-allocation-memberships | jq
+
+# Bump the user's SU sub-allocation.
+curl -s -X PUT $BASE/compute-allocation-memberships/$MEMBERSHIP_ID/allocation-amount \
+ -H 'Content-Type: application/json' \
+ -d '{"allocation_amount":75000}' | jq
+
+
+
# Record a usage diff against the allocation.
curl -s -X POST $BASE/compute-allocation-diffs \
-H 'Content-Type: application/json' \
diff --git a/internal/db/migrations/000010_compute_cluster_users.down.sql b/internal/db/migrations/000010_compute_cluster_users.down.sql
new file mode 100644
index 000000000..c207cd7c1
--- /dev/null
+++ b/internal/db/migrations/000010_compute_cluster_users.down.sql
@@ -0,0 +1,18 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+DROP TABLE IF EXISTS compute_cluster_users;
diff --git a/internal/db/migrations/000010_compute_cluster_users.up.sql b/internal/db/migrations/000010_compute_cluster_users.up.sql
new file mode 100644
index 000000000..284536208
--- /dev/null
+++ b/internal/db/migrations/000010_compute_cluster_users.up.sql
@@ -0,0 +1,31 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+CREATE TABLE IF NOT EXISTS compute_cluster_users
+(
+ id VARCHAR(255) NOT NULL,
+ compute_cluster_id VARCHAR(255) NOT NULL,
+ user_id VARCHAR(255) NOT NULL,
+ local_username VARCHAR(255) NOT NULL,
+ created_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),
+ updated_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6),
+ PRIMARY KEY (id),
+ UNIQUE KEY uq_compute_cluster_users_pair (compute_cluster_id, user_id),
+ KEY idx_compute_cluster_users_user (user_id),
+ CONSTRAINT fk_compute_cluster_users_cluster FOREIGN KEY (compute_cluster_id) REFERENCES compute_clusters (id) ON DELETE CASCADE,
+ CONSTRAINT fk_compute_cluster_users_user FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
+) ENGINE = InnoDB DEFAULT CHARSET = utf8mb4 COLLATE = utf8mb4_unicode_ci;
diff --git a/internal/server/server.go b/internal/server/server.go
index fb98b7320..4b3e20fd2 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -64,6 +64,14 @@ func (s *Server) routes() {
s.mux.HandleFunc("GET /compute-clusters", s.listComputeClusters)
s.mux.HandleFunc("GET /compute-clusters/{id}", s.getComputeCluster)
+ s.mux.HandleFunc("POST /compute-cluster-users", s.createComputeClusterUser)
+ s.mux.HandleFunc("GET /compute-cluster-users/{id}", s.getComputeClusterUser)
+ s.mux.HandleFunc("PUT /compute-cluster-users/{id}", s.updateComputeClusterUser)
+ s.mux.HandleFunc("DELETE /compute-cluster-users/{id}", s.deleteComputeClusterUser)
+ s.mux.HandleFunc("GET /compute-clusters/{id}/users", s.listComputeClusterUsersByCluster)
+ s.mux.HandleFunc("GET /compute-clusters/{id}/users/{userId}", s.getComputeClusterUserByPair)
+ s.mux.HandleFunc("GET /users/{id}/compute-cluster-users", s.listComputeClusterUsersByUser)
+
s.mux.HandleFunc("POST /compute-allocations", s.createComputeAllocation)
s.mux.HandleFunc("GET /compute-allocations/{id}", s.getComputeAllocation)
@@ -223,6 +231,78 @@ func (s *Server) listComputeClusters(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, clusters)
}
+func (s *Server) createComputeClusterUser(w http.ResponseWriter, r *http.Request) {
+ var cu models.ComputeClusterUser
+ if err := decodeJSON(r, &cu); err != nil {
+ writeError(w, http.StatusBadRequest, err)
+ return
+ }
+ created, err := s.svc.CreateComputeClusterUser(r.Context(), &cu)
+ if err != nil {
+ writeServiceError(w, err)
+ return
+ }
+ writeJSON(w, http.StatusCreated, created)
+}
+
+func (s *Server) getComputeClusterUser(w http.ResponseWriter, r *http.Request) {
+ cu, err := s.svc.GetComputeClusterUser(r.Context(), r.PathValue("id"))
+ if err != nil {
+ writeServiceError(w, err)
+ return
+ }
+ writeJSON(w, http.StatusOK, cu)
+}
+
+func (s *Server) updateComputeClusterUser(w http.ResponseWriter, r *http.Request) {
+ var cu models.ComputeClusterUser
+ if err := decodeJSON(r, &cu); err != nil {
+ writeError(w, http.StatusBadRequest, err)
+ return
+ }
+ cu.ID = r.PathValue("id")
+ if err := s.svc.UpdateComputeClusterUser(r.Context(), &cu); err != nil {
+ writeServiceError(w, err)
+ return
+ }
+ writeJSON(w, http.StatusOK, &cu)
+}
+
+func (s *Server) deleteComputeClusterUser(w http.ResponseWriter, r *http.Request) {
+ if err := s.svc.DeleteComputeClusterUser(r.Context(), r.PathValue("id")); err != nil {
+ writeServiceError(w, err)
+ return
+ }
+ w.WriteHeader(http.StatusNoContent)
+}
+
+func (s *Server) listComputeClusterUsersByCluster(w http.ResponseWriter, r *http.Request) {
+ users, err := s.svc.ListComputeClusterUsersByCluster(r.Context(), r.PathValue("id"))
+ if err != nil {
+ writeServiceError(w, err)
+ return
+ }
+ writeJSON(w, http.StatusOK, users)
+}
+
+func (s *Server) getComputeClusterUserByPair(w http.ResponseWriter, r *http.Request) {
+ cu, err := s.svc.GetComputeClusterUserByPair(r.Context(), r.PathValue("id"), r.PathValue("userId"))
+ if err != nil {
+ writeServiceError(w, err)
+ return
+ }
+ writeJSON(w, http.StatusOK, cu)
+}
+
+func (s *Server) listComputeClusterUsersByUser(w http.ResponseWriter, r *http.Request) {
+ users, err := s.svc.ListComputeClusterUsersByUser(r.Context(), r.PathValue("id"))
+ if err != nil {
+ writeServiceError(w, err)
+ return
+ }
+ writeJSON(w, http.StatusOK, users)
+}
+
func (s *Server) createComputeAllocation(w http.ResponseWriter, r *http.Request) {
var a models.ComputeAllocation
if err := decodeJSON(r, &a); err != nil {
diff --git a/internal/store/compute_cluster_user_store.go b/internal/store/compute_cluster_user_store.go
new file mode 100644
index 000000000..f5926ad1e
--- /dev/null
+++ b/internal/store/compute_cluster_user_store.go
@@ -0,0 +1,116 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package store
+
+import (
+ "context"
+ "database/sql"
+ "errors"
+
+ "github.com/jmoiron/sqlx"
+
+ "github.com/apache/airavata-custos/pkg/models"
+)
+
+type mysqlComputeClusterUserStore struct {
+ db *sqlx.DB
+}
+
+// NewComputeClusterUserStore returns a MySQL-backed ComputeClusterUserStore.
+func NewComputeClusterUserStore(db *sqlx.DB) ComputeClusterUserStore {
+ return &mysqlComputeClusterUserStore{db: db}
+}
+
+func (s *mysqlComputeClusterUserStore) FindByID(ctx context.Context, id string) (*models.ComputeClusterUser, error) {
+ var c models.ComputeClusterUser
+ err := s.db.GetContext(ctx, &c,
+ `SELECT id, compute_cluster_id, user_id, local_username
+ FROM compute_cluster_users WHERE id = ?`, id)
+ if err != nil {
+ if errors.Is(err, sql.ErrNoRows) {
+ return nil, nil
+ }
+ return nil, err
+ }
+ return &c, nil
+}
+
+func (s *mysqlComputeClusterUserStore) FindByPair(ctx context.Context, clusterID, userID string) (*models.ComputeClusterUser, error) {
+ var c models.ComputeClusterUser
+ err := s.db.GetContext(ctx, &c,
+ `SELECT id, compute_cluster_id, user_id, local_username
+ FROM compute_cluster_users
+ WHERE compute_cluster_id = ? AND user_id = ?`, clusterID, userID)
+ if err != nil {
+ if errors.Is(err, sql.ErrNoRows) {
+ return nil, nil
+ }
+ return nil, err
+ }
+ return &c, nil
+}
+
+func (s *mysqlComputeClusterUserStore) FindByCluster(ctx context.Context, clusterID string) ([]models.ComputeClusterUser, error) {
+ var users []models.ComputeClusterUser
+ err := s.db.SelectContext(ctx, &users,
+ `SELECT id, compute_cluster_id, user_id, local_username
+ FROM compute_cluster_users
+ WHERE compute_cluster_id = ?
+ ORDER BY local_username`, clusterID)
+ if err != nil {
+ return nil, err
+ }
+ return users, nil
+}
+
+func (s *mysqlComputeClusterUserStore) FindByUser(ctx context.Context, userID string) ([]models.ComputeClusterUser, error) {
+ var users []models.ComputeClusterUser
+ err := s.db.SelectContext(ctx, &users,
+ `SELECT id, compute_cluster_id, user_id, local_username
+ FROM compute_cluster_users
+ WHERE user_id = ?
+ ORDER BY compute_cluster_id`, userID)
+ if err != nil {
+ return nil, err
+ }
+ return users, nil
+}
+
+func (s *mysqlComputeClusterUserStore) Create(ctx context.Context, tx *sql.Tx, c *models.ComputeClusterUser) error {
+ _, err := tx.ExecContext(ctx,
+ `INSERT INTO compute_cluster_users (id, compute_cluster_id, user_id, local_username)
+ VALUES (?, ?, ?, ?)`,
+ c.ID, c.ComputeClusterID, c.UserID, c.LocalUsername)
+ return err
+}
+
+func (s *mysqlComputeClusterUserStore) Update(ctx context.Context, tx *sql.Tx, c *models.ComputeClusterUser) error {
+ _, err := tx.ExecContext(ctx,
+ `UPDATE compute_cluster_users
+ SET compute_cluster_id = ?,
+ user_id = ?,
+ local_username = ?
+ WHERE id = ?`,
+ c.ComputeClusterID, c.UserID, c.LocalUsername, c.ID)
+ return err
+}
+
+func (s *mysqlComputeClusterUserStore) Delete(ctx context.Context, tx *sql.Tx, id string) error {
+ _, err := tx.ExecContext(ctx, `DELETE FROM compute_cluster_users WHERE id = ?`, id)
+ return err
+}
diff --git a/internal/store/store.go b/internal/store/store.go
index ba6faa9f0..4dddec180 100644
--- a/internal/store/store.go
+++ b/internal/store/store.go
@@ -71,6 +71,25 @@ type ComputeClusterStore interface {
Delete(ctx context.Context, tx *sql.Tx, id string) error
}
+// ComputeClusterUserStore defines persistence operations for the mapping of
+// a Custos user to their local account on a compute cluster.
+type ComputeClusterUserStore interface {
+ // FindByID returns the mapping with the given ID, or nil if not found.
+ FindByID(ctx context.Context, id string) (*models.ComputeClusterUser, error)
+ // FindByPair returns the mapping for a (compute_cluster_id, user_id) pair, or nil if absent.
+ FindByPair(ctx context.Context, clusterID, userID string) (*models.ComputeClusterUser, error)
+ // FindByCluster returns every user mapping for the given compute cluster.
+ FindByCluster(ctx context.Context, clusterID string) ([]models.ComputeClusterUser, error)
+ // FindByUser returns every cluster mapping held by the given Custos user.
+ FindByUser(ctx context.Context, userID string) ([]models.ComputeClusterUser, error)
+ // Create inserts a new mapping within the provided transaction.
+ Create(ctx context.Context, tx *sql.Tx, c *models.ComputeClusterUser) error
+ // Update replaces mutable fields of an existing mapping within the provided transaction.
+ Update(ctx context.Context, tx *sql.Tx, c *models.ComputeClusterUser) error
+ // Delete removes a mapping by ID within the provided transaction.
+ Delete(ctx context.Context, tx *sql.Tx, id string) error
+}
+
// ProjectStore defines persistence operations for projects.
type ProjectStore interface {
// FindByID returns the project with the given ID, or nil if not found.
diff --git a/pkg/events/compute_cluster_user_subscribe.go b/pkg/events/compute_cluster_user_subscribe.go
new file mode 100644
index 000000000..5ba7a11a7
--- /dev/null
+++ b/pkg/events/compute_cluster_user_subscribe.go
@@ -0,0 +1,64 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package events
+
+import (
+ "log/slog"
+
+ "github.com/apache/airavata-custos/pkg/models"
+)
+
+// ComputeClusterUserHandler handles compute-cluster user lifecycle events
+// with a typed payload.
+type ComputeClusterUserHandler func(user models.ComputeClusterUser)
+
+// SubscribeComputeClusterUserCreated registers a typed handler invoked
+// whenever a compute_cluster_user::create event is published.
+func (b *Bus) SubscribeComputeClusterUserCreated(handler ComputeClusterUserHandler) {
+ b.subscribeComputeClusterUser(ComputeClusterUserCreateEvent, handler)
+}
+
+// SubscribeComputeClusterUserUpdated registers a typed handler invoked
+// whenever a compute_cluster_user::update event is published.
+func (b *Bus) SubscribeComputeClusterUserUpdated(handler ComputeClusterUserHandler) {
+ b.subscribeComputeClusterUser(ComputeClusterUserUpdateEvent, handler)
+}
+
+// SubscribeComputeClusterUserDeleted registers a typed handler invoked
+// whenever a compute_cluster_user::delete event is published.
+func (b *Bus) SubscribeComputeClusterUserDeleted(handler ComputeClusterUserHandler) {
+ b.subscribeComputeClusterUser(ComputeClusterUserDeleteEvent, handler)
+}
+
+func (b *Bus) subscribeComputeClusterUser(topic EventType, handler ComputeClusterUserHandler) {
+ b.Subscribe(topic, func(event Event, value interface{}) {
+ switch u := value.(type) {
+ case models.ComputeClusterUser:
+ handler(u)
+ case *models.ComputeClusterUser:
+ if u != nil {
+ handler(*u)
+ }
+ default:
+ slog.Warn("compute cluster user event payload has unexpected type",
+ "type", event.Type,
+ "got", value,
+ )
+ }
+ })
+}
diff --git a/pkg/events/types.go b/pkg/events/types.go
index 11c900e94..44758f0e4 100644
--- a/pkg/events/types.go
+++ b/pkg/events/types.go
@@ -54,6 +54,13 @@ const (
ComputeClusterDeleteEvent EventType = "compute_cluster::delete"
)
+// ComputeClusterUser lifecycle message types.
+const (
+ ComputeClusterUserCreateEvent EventType = "compute_cluster_user::create"
+ ComputeClusterUserUpdateEvent EventType = "compute_cluster_user::update"
+ ComputeClusterUserDeleteEvent EventType = "compute_cluster_user::delete"
+)
+
// ClusterAccount lifecycle message types.
const (
ClusterAccountCreateEvent EventType = "cluster_account::create"
diff --git a/pkg/models/allocation.go b/pkg/models/allocation.go
index 1037b1d7f..7813dcc2a 100644
--- a/pkg/models/allocation.go
+++ b/pkg/models/allocation.go
@@ -15,6 +15,13 @@ type ComputeCluster struct {
Name string `json:"name" db:"name"` // A human-readable name for the compute cluster, e.g., "Cluster A", "Cluster B", etc.
}
+type ComputeClusterUser struct {
+ ID string `json:"id" db:"id"`
+ ComputeClusterID string `json:"compute_cluster_id" db:"compute_cluster_id"`
+ UserID string `json:"user_id" db:"user_id"`
+ LocalUsername string `json:"local_username" db:"local_username"` // The username of the user on the compute cluster, which may be different from their Airavata Custos username.
+}
+
type ComputeAllocation struct {
ID string `json:"id" db:"id"`
ProjectID string `json:"project_id" db:"project_id"`
diff --git a/pkg/service/compute_cluster_user.go b/pkg/service/compute_cluster_user.go
new file mode 100644
index 000000000..3e8ed9868
--- /dev/null
+++ b/pkg/service/compute_cluster_user.go
@@ -0,0 +1,179 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package service
+
+import (
+ "context"
+ "database/sql"
+ "fmt"
+
+ "github.com/apache/airavata-custos/pkg/events"
+ "github.com/apache/airavata-custos/pkg/models"
+)
+
+// CreateComputeClusterUser persists a new compute-cluster user mapping. If
+// the ID is empty, a UUID is generated. The (possibly populated) record is
+// returned.
+func (s *Service) CreateComputeClusterUser(ctx context.Context, cu *models.ComputeClusterUser) (*models.ComputeClusterUser, error) {
+ if cu == nil {
+ return nil, fmt.Errorf("%w: compute cluster user is nil", ErrInvalidInput)
+ }
+ if cu.ComputeClusterID == "" {
+ return nil, fmt.Errorf("%w: compute_cluster_id is required", ErrInvalidInput)
+ }
+ if cu.UserID == "" {
+ return nil, fmt.Errorf("%w: user_id is required", ErrInvalidInput)
+ }
+ if cu.LocalUsername == "" {
+ return nil, fmt.Errorf("%w: local_username is required", ErrInvalidInput)
+ }
+ if cu.ID == "" {
+ cu.ID = newID()
+ }
+
+ if cluster, err := s.clusters.FindByID(ctx, cu.ComputeClusterID); err != nil {
+ return nil, fmt.Errorf("lookup compute cluster: %w", err)
+ } else if cluster == nil {
+ return nil, fmt.Errorf("%w: compute cluster %q not found", ErrInvalidInput, cu.ComputeClusterID)
+ }
+ if user, err := s.users.FindByID(ctx, cu.UserID); err != nil {
+ return nil, fmt.Errorf("lookup user: %w", err)
+ } else if user == nil {
+ return nil, fmt.Errorf("%w: user %q not found", ErrInvalidInput, cu.UserID)
+ }
+
+ if existing, err := s.clusterUsers.FindByPair(ctx, cu.ComputeClusterID, cu.UserID); err != nil {
+ return nil, fmt.Errorf("lookup compute cluster user by pair: %w", err)
+ } else if existing != nil {
+ return nil, fmt.Errorf("%w: user %q is already mapped on cluster %q",
+ ErrAlreadyExists, cu.UserID, cu.ComputeClusterID)
+ }
+
+ if err := s.inTx(ctx, func(tx *sql.Tx) error {
+ return s.clusterUsers.Create(ctx, tx, cu)
+ }); err != nil {
+ return nil, fmt.Errorf("create compute cluster user: %w", err)
+ }
+
+ s.eventBus.Publish(events.ComputeClusterUserCreateEvent, cu)
+ return cu, nil
+}
+
+// GetComputeClusterUser retrieves a compute-cluster user by its ID.
+func (s *Service) GetComputeClusterUser(ctx context.Context, id string) (*models.ComputeClusterUser, error) {
+ c, err := s.clusterUsers.FindByID(ctx, id)
+ if err != nil {
+ return nil, fmt.Errorf("get compute cluster user: %w", err)
+ }
+ if c == nil {
+ return nil, ErrNotFound
+ }
+ return c, nil
+}
+
+// GetComputeClusterUserByPair retrieves the compute-cluster user mapping for
+// the given (compute_cluster_id, user_id) pair.
+func (s *Service) GetComputeClusterUserByPair(ctx context.Context, clusterID, userID string) (*models.ComputeClusterUser, error) {
+ if clusterID == "" {
+ return nil, fmt.Errorf("%w: compute_cluster_id is required", ErrInvalidInput)
+ }
+ if userID == "" {
+ return nil, fmt.Errorf("%w: user_id is required", ErrInvalidInput)
+ }
+ c, err := s.clusterUsers.FindByPair(ctx, clusterID, userID)
+ if err != nil {
+ return nil, fmt.Errorf("get compute cluster user by pair: %w", err)
+ }
+ if c == nil {
+ return nil, ErrNotFound
+ }
+ return c, nil
+}
+
+// ListComputeClusterUsersByCluster returns every user mapping for the given
+// compute cluster, ordered by local username.
+func (s *Service) ListComputeClusterUsersByCluster(ctx context.Context, clusterID string) ([]models.ComputeClusterUser, error) {
+ if clusterID == "" {
+ return nil, fmt.Errorf("%w: compute_cluster_id is required", ErrInvalidInput)
+ }
+ users, err := s.clusterUsers.FindByCluster(ctx, clusterID)
+ if err != nil {
+ return nil, fmt.Errorf("list compute cluster users by cluster: %w", err)
+ }
+ return users, nil
+}
+
+// ListComputeClusterUsersByUser returns every cluster mapping held by the
+// given Custos user.
+func (s *Service) ListComputeClusterUsersByUser(ctx context.Context, userID string) ([]models.ComputeClusterUser, error) {
+ if userID == "" {
+ return nil, fmt.Errorf("%w: user_id is required", ErrInvalidInput)
+ }
+ users, err := s.clusterUsers.FindByUser(ctx, userID)
+ if err != nil {
+ return nil, fmt.Errorf("list compute cluster users by user: %w", err)
+ }
+ return users, nil
+}
+
+// UpdateComputeClusterUser persists changes to an existing compute-cluster
+// user mapping.
+func (s *Service) UpdateComputeClusterUser(ctx context.Context, cu *models.ComputeClusterUser) error {
+ if cu == nil || cu.ID == "" {
+ return fmt.Errorf("%w: compute cluster user id is required", ErrInvalidInput)
+ }
+ if cu.ComputeClusterID == "" {
+ return fmt.Errorf("%w: compute_cluster_id is required", ErrInvalidInput)
+ }
+ if cu.UserID == "" {
+ return fmt.Errorf("%w: user_id is required", ErrInvalidInput)
+ }
+ if cu.LocalUsername == "" {
+ return fmt.Errorf("%w: local_username is required", ErrInvalidInput)
+ }
+ if err := s.inTx(ctx, func(tx *sql.Tx) error {
+ return s.clusterUsers.Update(ctx, tx, cu)
+ }); err != nil {
+ return fmt.Errorf("update compute cluster user: %w", err)
+ }
+
+ s.eventBus.Publish(events.ComputeClusterUserUpdateEvent, cu)
+ return nil
+}
+
+// DeleteComputeClusterUser removes a compute-cluster user mapping by ID.
+func (s *Service) DeleteComputeClusterUser(ctx context.Context, id string) error {
+ if id == "" {
+ return fmt.Errorf("%w: compute cluster user id is required", ErrInvalidInput)
+ }
+ cu, err := s.clusterUsers.FindByID(ctx, id)
+ if err != nil {
+ return fmt.Errorf("lookup compute cluster user: %w", err)
+ }
+ if cu == nil {
+ return ErrNotFound
+ }
+ if err := s.inTx(ctx, func(tx *sql.Tx) error {
+ return s.clusterUsers.Delete(ctx, tx, id)
+ }); err != nil {
+ return fmt.Errorf("delete compute cluster user: %w", err)
+ }
+
+ s.eventBus.Publish(events.ComputeClusterUserDeleteEvent, cu)
+ return nil
+}
diff --git a/pkg/service/service.go b/pkg/service/service.go
index 935630ac5..e187ff82f 100644
--- a/pkg/service/service.go
+++ b/pkg/service/service.go
@@ -40,6 +40,7 @@ type Service struct {
users store.UserStore
projs store.ProjectStore
clusters store.ComputeClusterStore
+ clusterUsers store.ComputeClusterUserStore
allocs store.ComputeAllocationStore
resources store.ComputeAllocationResourceStore
resourceMappings store.ComputeAllocationResourceMappingStore
@@ -61,6 +62,7 @@ func New(database *sqlx.DB, eventBus *events.Bus) *Service {
users: store.NewUserStore(database),
projs: store.NewProjectStore(database),
clusters: store.NewComputeClusterStore(database),
+ clusterUsers: store.NewComputeClusterUserStore(database),
allocs: store.NewComputeAllocationStore(database),
resources: store.NewComputeAllocationResourceStore(database),
resourceMappings: store.NewComputeAllocationResourceMappingStore(database),
@@ -83,6 +85,7 @@ func NewWithStores(
users store.UserStore,
projs store.ProjectStore,
clusters store.ComputeClusterStore,
+ clusterUsers store.ComputeClusterUserStore,
allocs store.ComputeAllocationStore,
resources store.ComputeAllocationResourceStore,
resourceMappings store.ComputeAllocationResourceMappingStore,
@@ -100,6 +103,7 @@ func NewWithStores(
users: users,
projs: projs,
clusters: clusters,
+ clusterUsers: clusterUsers,
allocs: allocs,
resources: resources,
resourceMappings: resourceMappings,
(airavata-custos) 06/07: Updating account resource limits when the allocation resource mapping is created
Posted by di...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
DImuthuUpe pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/airavata-custos.git
commit 13ef537fd8b5a7438fc654f4ddd2084aac944b70
Author: DImuthuUpe <di...@gmail.com>
AuthorDate: Tue May 19 21:12:58 2026 -0400
Updating account resource limits when the allocation resource mapping is created
---
.../internal/subscribers/account.go | 60 ++++++++++++++++++++++
.../internal/subscribers/members.go | 4 +-
.../internal/subscribers/subscriber.go | 1 +
docs/API-Docs.md | 56 +++++++++++++++++---
...te_allocation_resource_mapping_amounts.down.sql | 20 ++++++++
...pute_allocation_resource_mapping_amounts.up.sql | 20 ++++++++
internal/server/server.go | 24 ++++++++-
.../compute_allocation_resource_mapping_store.go | 33 ++++++++++--
internal/store/store.go | 4 ++
pkg/models/allocation.go | 10 ++--
pkg/service/compute_allocation_resource_mapping.go | 50 ++++++++++++++++--
11 files changed, 262 insertions(+), 20 deletions(-)
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/account.go b/connectors/SLURM/Association-Mapper/internal/subscribers/account.go
index 2b872d54c..70a0f49e6 100644
--- a/connectors/SLURM/Association-Mapper/internal/subscribers/account.go
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/account.go
@@ -59,3 +59,63 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationDeletion(computeAcco
func (a *AssociationSubscriber) SubscribeToComputeAllocationUpdate(computeAccount models.ComputeAllocation) {
slog.Info("Received compute allocation update event", "account", computeAccount)
}
+
+func (a *AssociationSubscriber) SubscribeToComputeAllocationResourceMappingCreation(mapping models.ComputeAllocationResourceMapping) {
+ slog.Info("Received compute allocation resource mapping creation event", "mapping", mapping)
+
+ ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+ defer cancel()
+
+ allocation, err := a.coreService.GetComputeAllocation(ctx, mapping.ComputeAllocationID)
+ if err != nil {
+ slog.Error("Failed to get compute allocation for resource mapping creation", "error", err)
+ return
+ }
+
+ cluster, err := a.coreService.GetComputeCluster(ctx, allocation.ComputeClusterID)
+ if err != nil {
+ slog.Error("Failed to get compute cluster for resource mapping creation", "error", err)
+ return
+ }
+
+ resource, err := a.coreService.GetComputeAllocationResource(ctx, mapping.ComputeAllocationResourceID)
+ if err != nil {
+ slog.Error("Failed to get compute allocation resource for resource mapping creation", "error", err)
+ return
+ }
+
+ grpTres := []client.TRES{}
+
+ if mapping.ResourceAmount > 0 {
+ grpTres = append(grpTres, client.TRES{
+ Type: resource.ResourceType,
+ Count: mapping.ResourceAmount,
+ })
+ }
+
+ grpTresMins := []client.TRES{}
+ if mapping.ResourceTime > 0 {
+ grpTresMins = append(grpTresMins, client.TRES{
+ Type: resource.ResourceType,
+ Count: mapping.ResourceTime,
+ })
+ }
+
+ limits := client.AssocLimits{
+ GrpTRES: grpTres,
+ GrpTRESMins: grpTresMins,
+ }
+
+ association := client.Association{
+ Account: allocation.Name,
+ Cluster: cluster.Name,
+ Limits: limits,
+ }
+
+ err = a.slurmClient.UpsertAssociation(association)
+ if err != nil {
+ slog.Error("Failed to upsert association for membership resource override creation", "error", err)
+ } else {
+ slog.Info("Successfully upserted association for membership resource override creation", "association", association)
+ }
+}
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/members.go b/connectors/SLURM/Association-Mapper/internal/subscribers/members.go
index 687ecd8d5..dfcc4ee70 100644
--- a/connectors/SLURM/Association-Mapper/internal/subscribers/members.go
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/members.go
@@ -118,7 +118,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationMembershipResourceOv
if allocationResource.ResourceType == "GrpTRES" {
grpTres = append(grpTres, client.TRES{
- Type: allocationResource.Name,
+ Type: allocationResource.ResourceType,
Count: override.OverriddenResourceAmount, // override.OverriddenResourceAmount is the SU amount, but SLURM needs the actual resource amount (e.g., number of CPU hours), so we need to convert it using the rate for the resource
})
}
@@ -127,7 +127,7 @@ func (a *AssociationSubscriber) SubscribeToComputeAllocationMembershipResourceOv
if allocationResource.ResourceType == "GrpTRESMins" {
grpTresMins = append(grpTresMins, client.TRES{
- Type: allocationResource.Name,
+ Type: allocationResource.ResourceType,
Count: override.OverriddenResourceAmount,
})
}
diff --git a/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go b/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go
index 17b65cce8..308928a4b 100644
--- a/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go
+++ b/connectors/SLURM/Association-Mapper/internal/subscribers/subscriber.go
@@ -24,4 +24,5 @@ func (a *AssociationSubscriber) RegisterSubscribers() {
a.eventBus.SubscribeComputeAllocationUpdated(a.SubscribeToComputeAllocationUpdate)
a.eventBus.SubscribeComputeAllocationMembershipCreated(a.SubscribeToComputeAllocationMembershipCreation)
a.eventBus.SubscribeComputeAllocationMembershipResourceOverrideCreated(a.SubscribeToComputeAllocationMembershipResourceOverrideCreation)
+ a.eventBus.SubscribeComputeAllocationResourceMappingCreated(a.SubscribeToComputeAllocationResourceMappingCreation)
}
diff --git a/docs/API-Docs.md b/docs/API-Docs.md
index b571c8255..325a219bc 100644
--- a/docs/API-Docs.md
+++ b/docs/API-Docs.md
@@ -619,15 +619,21 @@ removed.
### `POST /compute-allocations/{id}/resources`
-Attach an existing resource to a compute allocation.
+Attach an existing resource to a compute allocation, recording the amount of
+the resource and the wall-clock time granted to the allocation.
**Path parameters:** `{id}` — the compute allocation ID.
**Required body fields:** `compute_allocation_resource_id`
+**Optional body fields:** `resource_amount` (int64, default `0`), `resource_time` (int64, default `0`). Both must be non-negative.
**Request**
```json
-{ "compute_allocation_resource_id": "c0a1b2c3-d4e5-46f7-8899-aabbccddeeff" }
+{
+ "compute_allocation_resource_id": "c0a1b2c3-d4e5-46f7-8899-aabbccddeeff",
+ "resource_amount": 24,
+ "resource_time": 1440
+}
```
**Response 201**
@@ -636,17 +642,55 @@ Attach an existing resource to a compute allocation.
{
"id": "7e1d2c3b-4a5f-4b6c-9d8e-0011223344ff",
"compute_allocation_id": "2f6a8c1d-3e4b-4a7d-8c91-aa12bb34cc56",
- "compute_allocation_resource_id": "c0a1b2c3-d4e5-46f7-8899-aabbccddeeff"
+ "compute_allocation_resource_id": "c0a1b2c3-d4e5-46f7-8899-aabbccddeeff",
+ "resource_amount": 24,
+ "resource_time": 1440
}
```
**Errors**
-- `400` — `compute_allocation_resource_id` missing, or either the allocation or the resource does not exist.
+- `400` — `compute_allocation_resource_id` missing, `resource_amount` / `resource_time` negative, or either the allocation or the resource does not exist.
- `409` — this resource is already attached to the allocation.
---
+### `PUT /compute-allocations/{id}/resources/{resourceId}`
+
+Update the `resource_amount` and `resource_time` recorded on an existing
+(allocation, resource) mapping.
+
+**Path parameters:** `{id}` — the compute allocation ID; `{resourceId}` — the compute allocation resource ID.
+**Required body fields:** `resource_amount`, `resource_time` (both non-negative int64).
+
+**Request**
+
+```json
+{
+ "resource_amount": 48,
+ "resource_time": 2880
+}
+```
+
+**Response 200**
+
+```json
+{
+ "id": "7e1d2c3b-4a5f-4b6c-9d8e-0011223344ff",
+ "compute_allocation_id": "2f6a8c1d-3e4b-4a7d-8c91-aa12bb34cc56",
+ "compute_allocation_resource_id": "c0a1b2c3-d4e5-46f7-8899-aabbccddeeff",
+ "resource_amount": 48,
+ "resource_time": 2880
+}
+```
+
+**Errors**
+
+- `400` — either id missing, or `resource_amount` / `resource_time` negative.
+- `404` — no such mapping exists.
+
+---
+
### `DELETE /compute-allocations/{id}/resources/{resourceId}`
Detach a resource from a compute allocation.
@@ -1276,7 +1320,7 @@ RES_ID=$(curl -s -X POST $BASE/compute-allocation-resources \
# Attach the resource to the allocation.
curl -s -X POST $BASE/compute-allocations/$ALLOC_ID/resources \
-H 'Content-Type: application/json' \
- -d "{\"compute_allocation_resource_id\":\"$RES_ID\"}" | jq
+ -d "{\"compute_allocation_resource_id\":\"$RES_ID\",\"resource_amount\":24,\"resource_time\":1440}" | jq
# Define a rate for the resource.
curl -s -X POST $BASE/compute-allocation-resource-rates \
@@ -1329,7 +1373,7 @@ OVERRIDE_ID=$(curl -s -X POST $BASE/compute-allocation-membership-resource-overr
-d "{
\"compute_allocation_membership_id\":\"$MEMBERSHIP_ID\",
\"compute_allocation_resource_id\":\"$RES_ID\",
- \"overridden_resource_amount\":10000
+ \"overridden_resource_amount\":10
}" | jq -r .id)
# Bump the override amount.
diff --git a/internal/db/migrations/000012_compute_allocation_resource_mapping_amounts.down.sql b/internal/db/migrations/000012_compute_allocation_resource_mapping_amounts.down.sql
new file mode 100644
index 000000000..b792b9027
--- /dev/null
+++ b/internal/db/migrations/000012_compute_allocation_resource_mapping_amounts.down.sql
@@ -0,0 +1,20 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+ALTER TABLE compute_allocation_resource_mappings
+ DROP COLUMN resource_time,
+ DROP COLUMN resource_amount;
diff --git a/internal/db/migrations/000012_compute_allocation_resource_mapping_amounts.up.sql b/internal/db/migrations/000012_compute_allocation_resource_mapping_amounts.up.sql
new file mode 100644
index 000000000..120491635
--- /dev/null
+++ b/internal/db/migrations/000012_compute_allocation_resource_mapping_amounts.up.sql
@@ -0,0 +1,20 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+ALTER TABLE compute_allocation_resource_mappings
+ ADD COLUMN resource_amount BIGINT NOT NULL DEFAULT 0 AFTER compute_allocation_resource_id,
+ ADD COLUMN resource_time BIGINT NOT NULL DEFAULT 0 AFTER resource_amount;
diff --git a/internal/server/server.go b/internal/server/server.go
index 960799fd8..6af71d2df 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -81,6 +81,7 @@ func (s *Server) routes() {
s.mux.HandleFunc("GET /compute-allocations/{id}/resources", s.listResourcesForAllocation)
s.mux.HandleFunc("POST /compute-allocations/{id}/resources", s.attachResourceToAllocation)
+ s.mux.HandleFunc("PUT /compute-allocations/{id}/resources/{resourceId}", s.updateAllocationResourceMapping)
s.mux.HandleFunc("DELETE /compute-allocations/{id}/resources/{resourceId}", s.detachResourceFromAllocation)
s.mux.HandleFunc("GET /compute-allocation-resources/{id}/allocations", s.listAllocationsForResource)
@@ -366,6 +367,8 @@ func (s *Server) listComputeAllocationResources(w http.ResponseWriter, r *http.R
type attachResourceRequest struct {
ComputeAllocationResourceID string `json:"compute_allocation_resource_id"`
+ ResourceAmount int64 `json:"resource_amount"`
+ ResourceTime int64 `json:"resource_time"`
}
func (s *Server) attachResourceToAllocation(w http.ResponseWriter, r *http.Request) {
@@ -374,7 +377,7 @@ func (s *Server) attachResourceToAllocation(w http.ResponseWriter, r *http.Reque
writeError(w, http.StatusBadRequest, err)
return
}
- mapping, err := s.svc.AttachResourceToAllocation(r.Context(), r.PathValue("id"), body.ComputeAllocationResourceID)
+ mapping, err := s.svc.AttachResourceToAllocation(r.Context(), r.PathValue("id"), body.ComputeAllocationResourceID, body.ResourceAmount, body.ResourceTime)
if err != nil {
writeServiceError(w, err)
return
@@ -382,6 +385,25 @@ func (s *Server) attachResourceToAllocation(w http.ResponseWriter, r *http.Reque
writeJSON(w, http.StatusCreated, mapping)
}
+type updateAllocationResourceMappingRequest struct {
+ ResourceAmount int64 `json:"resource_amount"`
+ ResourceTime int64 `json:"resource_time"`
+}
+
+func (s *Server) updateAllocationResourceMapping(w http.ResponseWriter, r *http.Request) {
+ var body updateAllocationResourceMappingRequest
+ if err := decodeJSON(r, &body); err != nil {
+ writeError(w, http.StatusBadRequest, err)
+ return
+ }
+ mapping, err := s.svc.UpdateAllocationResourceMapping(r.Context(), r.PathValue("id"), r.PathValue("resourceId"), body.ResourceAmount, body.ResourceTime)
+ if err != nil {
+ writeServiceError(w, err)
+ return
+ }
+ writeJSON(w, http.StatusOK, mapping)
+}
+
func (s *Server) detachResourceFromAllocation(w http.ResponseWriter, r *http.Request) {
if err := s.svc.DetachResourceFromAllocation(r.Context(), r.PathValue("id"), r.PathValue("resourceId")); err != nil {
writeServiceError(w, err)
diff --git a/internal/store/compute_allocation_resource_mapping_store.go b/internal/store/compute_allocation_resource_mapping_store.go
index ca89bbf06..9f05bf70b 100644
--- a/internal/store/compute_allocation_resource_mapping_store.go
+++ b/internal/store/compute_allocation_resource_mapping_store.go
@@ -37,10 +37,25 @@ func NewComputeAllocationResourceMappingStore(db *sqlx.DB) ComputeAllocationReso
return &mysqlComputeAllocationResourceMappingStore{db: db}
}
+func (s *mysqlComputeAllocationResourceMappingStore) FindByID(ctx context.Context, id string) (*models.ComputeAllocationResourceMapping, error) {
+ var m models.ComputeAllocationResourceMapping
+ err := s.db.GetContext(ctx, &m,
+ `SELECT id, compute_allocation_id, compute_allocation_resource_id, resource_amount, resource_time
+ FROM compute_allocation_resource_mappings
+ WHERE id = ?`, id)
+ if err != nil {
+ if errors.Is(err, sql.ErrNoRows) {
+ return nil, nil
+ }
+ return nil, err
+ }
+ return &m, nil
+}
+
func (s *mysqlComputeAllocationResourceMappingStore) FindByPair(ctx context.Context, allocationID, resourceID string) (*models.ComputeAllocationResourceMapping, error) {
var m models.ComputeAllocationResourceMapping
err := s.db.GetContext(ctx, &m,
- `SELECT id, compute_allocation_id, compute_allocation_resource_id
+ `SELECT id, compute_allocation_id, compute_allocation_resource_id, resource_amount, resource_time
FROM compute_allocation_resource_mappings
WHERE compute_allocation_id = ? AND compute_allocation_resource_id = ?`,
allocationID, resourceID)
@@ -87,9 +102,19 @@ func (s *mysqlComputeAllocationResourceMappingStore) FindAllocationsByResource(c
func (s *mysqlComputeAllocationResourceMappingStore) Create(ctx context.Context, tx *sql.Tx, m *models.ComputeAllocationResourceMapping) error {
_, err := tx.ExecContext(ctx,
`INSERT INTO compute_allocation_resource_mappings
- (id, compute_allocation_id, compute_allocation_resource_id)
- VALUES (?, ?, ?)`,
- m.ID, m.ComputeAllocationID, m.ComputeAllocationResourceID)
+ (id, compute_allocation_id, compute_allocation_resource_id, resource_amount, resource_time)
+ VALUES (?, ?, ?, ?, ?)`,
+ m.ID, m.ComputeAllocationID, m.ComputeAllocationResourceID, m.ResourceAmount, m.ResourceTime)
+ return err
+}
+
+func (s *mysqlComputeAllocationResourceMappingStore) Update(ctx context.Context, tx *sql.Tx, m *models.ComputeAllocationResourceMapping) error {
+ _, err := tx.ExecContext(ctx,
+ `UPDATE compute_allocation_resource_mappings
+ SET resource_amount = ?,
+ resource_time = ?
+ WHERE id = ?`,
+ m.ResourceAmount, m.ResourceTime, m.ID)
return err
}
diff --git a/internal/store/store.go b/internal/store/store.go
index 8066d87e8..b066df428 100644
--- a/internal/store/store.go
+++ b/internal/store/store.go
@@ -140,6 +140,8 @@ type ComputeAllocationResourceStore interface {
// ComputeAllocationResourceMappingStore defines persistence operations for
// the join table linking compute allocations and compute allocation resources.
type ComputeAllocationResourceMappingStore interface {
+ // FindByID returns the mapping with the given ID, or nil if it does not exist.
+ FindByID(ctx context.Context, id string) (*models.ComputeAllocationResourceMapping, error)
// FindByPair returns the mapping for a (allocation, resource) pair, or nil if absent.
FindByPair(ctx context.Context, allocationID, resourceID string) (*models.ComputeAllocationResourceMapping, error)
// FindResourcesByAllocation returns every resource attached to the given allocation.
@@ -148,6 +150,8 @@ type ComputeAllocationResourceMappingStore interface {
FindAllocationsByResource(ctx context.Context, resourceID string) ([]models.ComputeAllocation, error)
// Create inserts a new mapping within the provided transaction.
Create(ctx context.Context, tx *sql.Tx, m *models.ComputeAllocationResourceMapping) error
+ // Update replaces mutable fields of an existing mapping within the provided transaction.
+ Update(ctx context.Context, tx *sql.Tx, m *models.ComputeAllocationResourceMapping) error
// DeleteByPair removes the mapping for a (allocation, resource) pair within the provided transaction.
DeleteByPair(ctx context.Context, tx *sql.Tx, allocationID, resourceID string) error
}
diff --git a/pkg/models/allocation.go b/pkg/models/allocation.go
index ca95595b5..582bce606 100644
--- a/pkg/models/allocation.go
+++ b/pkg/models/allocation.go
@@ -33,17 +33,21 @@ type ComputeAllocation struct {
EndTime time.Time `json:"end_time" db:"end_time"`
}
+// Typically store the a paritition information
type ComputeAllocationResource struct {
ID string `json:"id" db:"id"`
- Name string `json:"name" db:"name"` // resource / partition name, e.g., "cpu", "gpu", etc.
- ResourceType string `json:"resource_type" db:"resource_type"` // GrpTRES, GrpTRESMins
- ResourceAmount int64 `json:"resource_amount" db:"resource_amount"` // Number of CPUs, GPUs, time in minutes, or other unit depending on the resource type.
+ Name string `json:"name" db:"name"` // resource / partition name, e.g., "cpu-01", "gpu-01", "gpu-interactive", etc.
+ ResourceType string `json:"resource_type" db:"resource_type"` // cpu, gpu
+ ResourceAmount int64 `json:"resource_amount" db:"resource_amount"` // Number of CPUs, GPUs.
}
+// Store the association amount for a parition and allocation
type ComputeAllocationResourceMapping struct {
ID string `json:"id" db:"id"`
ComputeAllocationID string `json:"compute_allocation_id" db:"compute_allocation_id"`
ComputeAllocationResourceID string `json:"compute_allocation_resource_id" db:"compute_allocation_resource_id"`
+ ResourceAmount int64 `json:"resource_amount" db:"resource_amount"` // Amount of the resource allocated to this allocation (e.g., number of CPUs, GPUs).
+ ResourceTime int64 `json:"resource_time" db:"resource_time"` // Wall-clock time in minutes that the allocated amount is granted for.
}
type ComputeAllocationResourceRate struct {
diff --git a/pkg/service/compute_allocation_resource_mapping.go b/pkg/service/compute_allocation_resource_mapping.go
index d0c96014c..ce5f4ef18 100644
--- a/pkg/service/compute_allocation_resource_mapping.go
+++ b/pkg/service/compute_allocation_resource_mapping.go
@@ -27,16 +27,23 @@ import (
)
// AttachResourceToAllocation links a compute allocation resource to a compute
-// allocation. Both entities must already exist. The link is idempotent — if
-// the same (allocation, resource) pair is already linked, ErrAlreadyExists is
-// returned.
-func (s *Service) AttachResourceToAllocation(ctx context.Context, allocationID, resourceID string) (*models.ComputeAllocationResourceMapping, error) {
+// allocation, recording the resource amount and wall-clock time granted to
+// the allocation. Both entities must already exist. The link is idempotent —
+// if the same (allocation, resource) pair is already linked, ErrAlreadyExists
+// is returned.
+func (s *Service) AttachResourceToAllocation(ctx context.Context, allocationID, resourceID string, resourceAmount, resourceTime int64) (*models.ComputeAllocationResourceMapping, error) {
if allocationID == "" {
return nil, fmt.Errorf("%w: compute_allocation_id is required", ErrInvalidInput)
}
if resourceID == "" {
return nil, fmt.Errorf("%w: compute_allocation_resource_id is required", ErrInvalidInput)
}
+ if resourceAmount < 0 {
+ return nil, fmt.Errorf("%w: resource_amount must be non-negative", ErrInvalidInput)
+ }
+ if resourceTime < 0 {
+ return nil, fmt.Errorf("%w: resource_time must be non-negative", ErrInvalidInput)
+ }
if alloc, err := s.allocs.FindByID(ctx, allocationID); err != nil {
return nil, fmt.Errorf("lookup compute allocation: %w", err)
@@ -60,6 +67,8 @@ func (s *Service) AttachResourceToAllocation(ctx context.Context, allocationID,
ID: newID(),
ComputeAllocationID: allocationID,
ComputeAllocationResourceID: resourceID,
+ ResourceAmount: resourceAmount,
+ ResourceTime: resourceTime,
}
if err := s.inTx(ctx, func(tx *sql.Tx) error {
return s.resourceMappings.Create(ctx, tx, mapping)
@@ -71,6 +80,39 @@ func (s *Service) AttachResourceToAllocation(ctx context.Context, allocationID,
return mapping, nil
}
+// UpdateAllocationResourceMapping updates the resource_amount and
+// resource_time recorded against an existing (allocation, resource) mapping.
+func (s *Service) UpdateAllocationResourceMapping(ctx context.Context, allocationID, resourceID string, resourceAmount, resourceTime int64) (*models.ComputeAllocationResourceMapping, error) {
+ if allocationID == "" || resourceID == "" {
+ return nil, fmt.Errorf("%w: allocation and resource ids are required", ErrInvalidInput)
+ }
+ if resourceAmount < 0 {
+ return nil, fmt.Errorf("%w: resource_amount must be non-negative", ErrInvalidInput)
+ }
+ if resourceTime < 0 {
+ return nil, fmt.Errorf("%w: resource_time must be non-negative", ErrInvalidInput)
+ }
+
+ existing, err := s.resourceMappings.FindByPair(ctx, allocationID, resourceID)
+ if err != nil {
+ return nil, fmt.Errorf("lookup mapping: %w", err)
+ }
+ if existing == nil {
+ return nil, ErrNotFound
+ }
+
+ existing.ResourceAmount = resourceAmount
+ existing.ResourceTime = resourceTime
+ if err := s.inTx(ctx, func(tx *sql.Tx) error {
+ return s.resourceMappings.Update(ctx, tx, existing)
+ }); err != nil {
+ return nil, fmt.Errorf("update allocation resource mapping: %w", err)
+ }
+
+ s.eventBus.Publish(events.ComputeAllocationResourceMappingUpdateEvent, existing)
+ return existing, nil
+}
+
// DetachResourceFromAllocation removes the link between a compute allocation
// and a compute allocation resource. Returns ErrNotFound when no such mapping
// exists.