Increase the RLIMIT_MEMLOCK when TPUProxy is enabled.

PiperOrigin-RevId: 619045692
This commit is contained in:
Jing Chen
2024-03-25 20:30:08 -07:00
committed by gVisor bot
parent f94df6d1bf
commit cc37e536cb
2 changed files with 11 additions and 7 deletions
+6 -2
View File
@@ -113,12 +113,16 @@ func (d *defs) initDefaults() error {
return nil
}
func createLimitSet(spec *specs.Spec) (*limits.LimitSet, error) {
func createLimitSet(spec *specs.Spec, enableTPUProxy bool) (*limits.LimitSet, error) {
ls, err := defaults.get()
if err != nil {
return nil, err
}
// Set RLIMIT_MEMLOCK's default value to unlimited when TPUProxy is enabled.
// The value will be overwritten if the exact rlimit is provided.
if enableTPUProxy {
ls.SetUnchecked(limits.MemoryLocked, limits.Limit{Cur: limits.Infinity, Max: limits.Infinity})
}
// Then apply overwrites on top of defaults.
for _, rl := range spec.Process.Rlimits {
lt, ok := limits.FromLinuxResourceName[rl.Type]
+5 -5
View File
@@ -514,7 +514,7 @@ func New(args Args) (*Loader, error) {
dogOpts.TaskTimeoutAction = args.Conf.WatchdogAction
dog := watchdog.New(k, dogOpts)
procArgs, err := createProcessArgs(args.ID, args.Spec, creds, k, k.RootPIDNamespace())
procArgs, err := createProcessArgs(args.ID, args.Spec, args.Conf, creds, k, k.RootPIDNamespace())
if err != nil {
return nil, fmt.Errorf("creating init process for root container: %w", err)
}
@@ -582,9 +582,9 @@ func New(args Args) (*Loader, error) {
}
// createProcessArgs creates args that can be used with kernel.CreateProcess.
func createProcessArgs(id string, spec *specs.Spec, creds *auth.Credentials, k *kernel.Kernel, pidns *kernel.PIDNamespace) (kernel.CreateProcessArgs, error) {
func createProcessArgs(id string, spec *specs.Spec, conf *config.Config, creds *auth.Credentials, k *kernel.Kernel, pidns *kernel.PIDNamespace) (kernel.CreateProcessArgs, error) {
// Create initial limits.
ls, err := createLimitSet(spec)
ls, err := createLimitSet(spec, specutils.TPUProxyIsEnabled(spec, conf))
if err != nil {
return kernel.CreateProcessArgs{}, fmt.Errorf("creating limits: %w", err)
}
@@ -899,7 +899,7 @@ func (l *Loader) startSubcontainer(spec *specs.Spec, conf *config.Config, cid st
nvidiaDriverVersion: l.root.nvidiaDriverVersion,
}
var err error
info.procArgs, err = createProcessArgs(cid, spec, creds, l.k, pidns)
info.procArgs, err = createProcessArgs(cid, spec, conf, creds, l.k, pidns)
if err != nil {
return fmt.Errorf("creating new process: %w", err)
}
@@ -1196,7 +1196,7 @@ func (l *Loader) executeAsync(args *control.ExecArgs) (kernel.ThreadID, error) {
}
args.PIDNamespace = tg.PIDNamespace()
args.Limits, err = createLimitSet(l.root.spec)
args.Limits, err = createLimitSet(l.root.spec, specutils.TPUProxyIsEnabled(l.root.spec, l.root.conf))
if err != nil {
return 0, fmt.Errorf("creating limits: %w", err)
}