diff --git a/cmd/fetch.go b/cmd/fetch.go
index 3d99fb0..402b40c 100644
--- a/cmd/fetch.go
+++ b/cmd/fetch.go
@@ -106,6 +106,21 @@ func init() {
// fetch quota
fetchCmd.AddCommand(fetchQuotaCmd)
+
+ // fetch capacity
+ fetchCmd.AddCommand(fetchAvailCapacityCmd)
+
+ // Hijack help function to hide unnecessary global flags
+ fetchAvailCapacityCmd.SetHelpFunc(func(cmd *cobra.Command, s []string) {
+ if cmd.HasInheritedFlags() {
+ cmd.InheritedFlags().VisitAll(func(f *pflag.Flag) {
+ if f.Name != "logLevel" {
+ f.Hidden = true
+ }
+ })
+ }
+ help(cmd, s)
+ })
}
var fetchCmd = &cobra.Command{
@@ -183,6 +198,14 @@ var fetchQuotaCmd = &cobra.Command{
Run: fetchQuota,
}
+var fetchAvailCapacityCmd = &cobra.Command{
+ Use: "capacity",
+ PreRun: setConfig,
+ Short: "Fetch capacity report",
+ Long: `This command will show detailed capacity report of the cluster`,
+ Run: fetchAvailCapacity,
+}
+
func fetchTasksConfig(cmd *cobra.Command, args []string) {
log.Infof("Fetching job configuration for [%s/%s/%s] \n", *env, *role, *name)
@@ -416,3 +439,40 @@ func fetchQuota(cmd *cobra.Command, args []string) {
}
}
}
+
+//fetchAvailCapacity reports free capacity in details
+func fetchAvailCapacity(cmd *cobra.Command, args []string) {
+ log.Infof("Fetching available capacity from %s/offers\n", client.GetSchedulerURL())
+
+ report, err := client.AvailOfferReport()
+ if err != nil {
+ log.Fatalf("error: %+v\n", err)
+ }
+
+ // convert report to user-friendly structure
+ capacity := map[string]map[string]map[string]int64{}
+ for g, gv := range report {
+ if _, ok := capacity[g]; !ok {
+ capacity[g] = map[string]map[string]int64{}
+ }
+
+ for r, rc := range gv {
+ if _, ok := capacity[g][r]; !ok {
+ capacity[g][r] = map[string]int64{}
+ }
+
+ for v, c := range rc {
+ capacity[g][r][fmt.Sprint(v)] = c
+ }
+ }
+ }
+
+ if toJson {
+ fmt.Println(internal.ToJSON(capacity))
+ if err != nil {
+ log.Fatalf("error: %+v\n", err)
+ }
+ } else {
+ fmt.Println(capacity)
+ }
+}
diff --git a/cmd/root.go b/cmd/root.go
index 8e5532a..598a9fb 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -29,7 +29,8 @@ import (
var username, password, zkAddr, schedAddr string
var env, role, name = new(string), new(string), new(string)
-var ram, disk int64
+var dedicated string
+var ram, disk, gpu, port int64
var cpu float64
var client *realis.Client
var skipCertVerification bool
diff --git a/cmd/simulate.go b/cmd/simulate.go
new file mode 100644
index 0000000..a292693
--- /dev/null
+++ b/cmd/simulate.go
@@ -0,0 +1,61 @@
+/**
+ * Licensed 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 cmd
+
+import (
+ "fmt"
+
+ "github.com/aurora-scheduler/australis/internal"
+ "github.com/spf13/cobra"
+)
+
+func init() {
+ rootCmd.AddCommand(simulateCmd)
+
+ simulateCmd.AddCommand(fitCmd)
+}
+
+var simulateCmd = &cobra.Command{
+ Use: "simulate",
+ Short: "Simulate some work based on the current cluster condition, and return the output",
+}
+
+var fitCmd = &cobra.Command{
+ Use: "fit",
+ Short: "Compute how many tasks can we fit to a cluster",
+ Run: fit,
+ Args: cobra.RangeArgs(1, 2),
+}
+
+func fit(cmd *cobra.Command, args []string) {
+ log.Infof("Compute how many tasks can be fit in the remaining cluster capacity")
+
+ taskConfig, err := internal.UnmarshalTaskConfig(args[0])
+ if err != nil {
+ log.Fatalln(err)
+ }
+
+ offers, err := client.Offers()
+ if err != nil {
+ log.Fatal("error: %+v", err)
+ }
+
+ numTasks, err := client.FitTasks(taskConfig, offers)
+ if err != nil {
+ log.Fatal("error: %+v", err)
+ }
+
+ fmt.Println(numTasks)
+}
diff --git a/go.mod b/go.mod
index 8bfffbe..a947918 100644
--- a/go.mod
+++ b/go.mod
@@ -3,7 +3,9 @@ module github.com/aurora-scheduler/australis
go 1.15
require (
- github.com/aurora-scheduler/gorealis/v2 v2.26.0
+ github.com/andygrunwald/megos v0.0.0-20210622170559-e9ff1cac83c5
+ github.com/aurora-scheduler/gorealis/v2 v2.28.0
+ github.com/gizak/termui/v3 v3.1.0
github.com/pkg/errors v0.9.1
github.com/sirupsen/logrus v1.6.0
github.com/spf13/cobra v1.0.0
@@ -13,4 +15,4 @@ require (
gopkg.in/yaml.v2 v2.2.8
)
-replace github.com/apache/thrift v0.13.0 => github.com/ridv/thrift v0.13.2
+replace github.com/apache/thrift v0.13.0 => github.com/ridv/thrift v0.13.2
\ No newline at end of file
diff --git a/internal/job.go b/internal/job.go
index 181ea17..2f5e09b 100644
--- a/internal/job.go
+++ b/internal/job.go
@@ -65,6 +65,8 @@ type Job struct {
CPU float64 `yaml:"cpu"`
RAM int64 `yaml:"ram"`
Disk int64 `yaml:"disk"`
+ Port int64 `yaml:"port"`
+ GPU int64 `yaml:"gpu"`
Executor Executor `yaml:"executor"`
Instances int32 `yaml:"instances"`
MaxFailures int32 `yaml:"maxFailures"`
@@ -90,6 +92,8 @@ func (j *Job) ToRealis() (*realis.AuroraJob, error) {
CPU(j.CPU).
RAM(j.RAM).
Disk(j.Disk).
+ AddPorts(int(j.Port)).
+ GPU(j.GPU).
IsService(j.Service).
Tier(j.Tier).
Priority(j.Priority).
diff --git a/internal/util.go b/internal/util.go
index cc81cdb..dbf56be 100644
--- a/internal/util.go
+++ b/internal/util.go
@@ -118,6 +118,26 @@ func UnmarshalJob(filename string) (Job, error) {
return job, nil
}
+func UnmarshalTaskConfig(filename string) (*aurora.TaskConfig, error) {
+ if jobsFile, err := os.Open(filename); err != nil {
+ return nil, errors.Wrap(err, "unable to read the task config file")
+ } else {
+ job := Job{}
+
+ if err := yaml.NewDecoder(jobsFile).Decode(&job); err != nil {
+ return nil, errors.Wrap(err, "unable to parse task config file")
+ }
+
+ if auroraJob, err := job.ToRealis(); err != nil {
+ return nil, errors.Wrap(err, "unable to parse task config file")
+ } else {
+ return auroraJob.JobConfig().TaskConfig, nil
+ }
+ }
+
+ return nil, nil
+}
+
func UnmarshalUpdate(filename string) (UpdateJob, error) {
updateJob := UpdateJob{}
diff --git a/test/task_config.yaml b/test/task_config.yaml
new file mode 100644
index 0000000..addb8c1
--- /dev/null
+++ b/test/task_config.yaml
@@ -0,0 +1,20 @@
+environment: "prod"
+role: "vagrant"
+name: "hello_world"
+cpu: 0.09
+ram: 64
+disk: 128
+valueConstraints:
+ - name: "dedicated"
+ values:
+ - "vagrant/bar"
+limitConstraints:
+ - name: "host"
+ limit: 1
+ - name: "zone"
+ limit: 2
+thermos:
+ - name: "bootstrap"
+ cmd: "echo bootstrapping"
+ - name: "hello_gorealis"
+ cmd: "while true; do echo hello world from gorealis; sleep 10; done"