fetch capacity and simulate task fitting

cmd/fetch.go

@@ -106,6 +106,21 @@ func init() {
 
 	// fetch quota
 	fetchCmd.AddCommand(fetchQuotaCmd)
+
+	// fetch capacity
+	fetchCmd.AddCommand(fetchCapacityCmd)
+
+	// Hijack help function to hide unnecessary global flags
+	fetchCapacityCmd.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 fetchCapacityCmd = &cobra.Command{
+	Use:    "capacity",
+	PreRun: setConfig,
+	Short:  "Fetch capacity report",
+	Long:   `This command will show detailed capacity report of the cluster`,
+	Run:    fetchCapacity,
+}
+
 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) {
 		}
 	}
 }
+
+//fetchCapacity reports free capacity in details
+func fetchCapacity(cmd *cobra.Command, args []string) {
+	log.Infof("Fetching 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)
+	}
+}

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

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)
+}

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).

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{}

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"