Support crun/calloc

PKUHPC · Feb 23, 2025 · 50eb401 · 50eb401
1 parent 309ed61
commit 50eb401
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 1 deletion.
diff --git a/internal/calloc/CmdArgParser.go b/internal/calloc/CmdArgParser.go
@@ -45,6 +45,8 @@ var (
 	FlagConfigFilePath string
 	FlagDebugLevel     string
 
+	FlagReservation string
+
 	RootCmd = &cobra.Command{
 		Use:     "calloc",
 		Short:   "Allocate resource and create terminal",
@@ -87,4 +89,5 @@ func init() {
 	RootCmd.Flags().StringVarP(&FlagExcludes, "exclude", "x", "", "Exclude specific nodes from allocating (commas separated list)")
 	RootCmd.Flags().BoolVar(&FlagGetUserEnv, "get-user-env", false, "Load login environment variables of the user")
 	RootCmd.Flags().StringVar(&FlagExport, "export", "", "Propagate environment variables")
+	RootCmd.Flags().StringVarP(&FlagReservation, "reservation", "r", "", "Use reserved resources")
 }
diff --git a/internal/calloc/calloc.go b/internal/calloc/calloc.go
@@ -444,6 +444,9 @@ func MainCalloc(cmd *cobra.Command, args []string) util.CraneCmdError {
 	if FlagExport != "" {
 		task.Env["CRANE_EXPORT_ENV"] = FlagExport
 	}
+	if FlagReservation != "" {
+		task.Reservation = FlagReservation
+	}
 
 	// Set total limit of cpu cores
 	task.Resources.AllocatableRes.CpuCoreLimit = task.CpusPerTask * float64(task.NtasksPerNode)

diff --git a/internal/cbatch/CmdArgParser.go b/internal/cbatch/CmdArgParser.go
@@ -158,5 +158,5 @@ func init() {
 	RootCmd.Flags().StringVar(&FlagMailType, "mail-type", "", "Notify user by mail when certain events occur, supported values: NONE, BEGIN, END, FAIL, ALL (default is NONE)")
 	RootCmd.Flags().StringVar(&FlagMailUser, "mail-user", "", "Mail address of the notification receiver")
 	RootCmd.Flags().BoolVar(&FlagJson, "json", false, "Output in JSON format")
-	RootCmd.Flags().StringVarP(&FlagReservation, "reservation", "r", "", "Allow to user reservation")
+	RootCmd.Flags().StringVarP(&FlagReservation, "reservation", "r", "", "Use reserved resources")
 }
diff --git a/internal/crun/CmdArgParser.go b/internal/crun/CmdArgParser.go
@@ -59,6 +59,8 @@ var (
 
 	FlagConfigFilePath string
 	FlagDebugLevel     string
+
+	FlagReservation string
 )
 
 func ParseCmdArgs() {
@@ -89,4 +91,5 @@ func init() {
 	RootCmd.Flags().StringVar(&FlagExport, "export", "", "Propagate environment variables")
 	RootCmd.Flags().BoolVar(&FlagPty, "pty", false, "Run with a pseudo-terminal")
 	RootCmd.Flags().BoolVar(&FlagX11, "x11", false, "Run with X11 forwarding (insecure)")
+	RootCmd.Flags().StringVarP(&FlagReservation, "reservation", "r", "", "Use reserved resources")
 }
diff --git a/internal/crun/crun.go b/internal/crun/crun.go
@@ -683,6 +683,9 @@ func MainCrun(cmd *cobra.Command, args []string) util.CraneCmdError {
 	if FlagExport != "" {
 		task.Env["CRANE_EXPORT_ENV"] = FlagExport
 	}
+	if FlagReservation != "" {
+		task.Reservation = FlagReservation
+	}
 
 	// Set total limit of cpu cores
 	task.Resources.AllocatableRes.CpuCoreLimit = task.CpusPerTask * float64(task.NtasksPerNode)