Repository: kubeflow/training-operator (language: python)

Issue:
Infinitely looping init-pytorch container may eventually exceed its memory limit
The `init-pytorch` container generally uses very little memory and is expected to run for a relatively short time.  However, in a failure condition where the master pod will never be reachable, this container will loop forever ... until, depending on your host OS, it may eventually terminate and result in a worker pod status of `Init:OOMKilled` -- which is a bit confusing.

The root cause is that this simple loop can slowly increases in its memory consumption, possibly due to a memory leak or perhaps more likely due to OS-level memory fragmentation, and may eventually exceed the hard-coded `20Mi` memory limit (see line 42 below).

https://github.com/kubeflow/training-operator/blob/54a408fabfbe58d895aa0b6ec055b43f59a64da0/pkg/controller.v1/pytorch/initcontainer.go#L36-L46

One approach would be to simply increase the memory limit.  However, it might also be nice if that loop could be exited with failure status after a certain timeout interval has elapsed.


Discussion / hints:
@arlofaria Thanks for creating the issue. Can you add a fix if you can? You can keep the timeout interval to be a configurable parameter
Hi @johnugeorge, thanks for looking at this!

Unfortunately, my current employer's policies make it very cumbersome for me to contribute to open-source projects -- and I'm not particularly familiar with this codebase or how to properly test it either.

However, I would think that a simple fix would be to modify `pkg/config/config.go` to add something like `PyTorchInitContainerMaxTries` ... and then adjust the logic of the one-line command in `pytorch/initcontainer.go` to use something like a finite loop (e.g. `for i in $(seq {{ .InitContainerMaxTries }}); do ...`) and then make sure to break out of the loop with a successful `exit 0` if the master node is found, or otherwise `exit 1` after the finite loop is completed.

You could also do something like configuring a `PyTorchInitContainerTimeout`, which would have somewhat more complex logic ... unless you changed the loop to `sleep 1` instead of `sleep 2`, and then it'd be approximately the same as `PyTorchInitContainerMaxTries`.

Hope this helps, and sorry that I can't contribute more directly 😞 
/kind feature

I just recently learned about kubeflow and am very interested in it. I'm looking for opportunity to contribute and this looks like a good one to start. I'd like to have a try.
Hi, @kidddddddddddddddddddddd. Welcome to the kubeflow community. 

Once you are ready to work on this issue, feel free to assign yourself to this issue with `/assign`, the same as k/k.
/assign 
Will start working on this weekend.

Fix patch:
diff --git a/cmd/training-operator.v1/main.go b/cmd/training-operator.v1/main.go
index 4b09586cdb..164a019c62 100644
--- a/cmd/training-operator.v1/main.go
+++ b/cmd/training-operator.v1/main.go
@@ -85,6 +85,8 @@ func main() {
 		config.PyTorchInitContainerImageDefault, "The image for pytorch init container")
 	flag.StringVar(&config.Config.PyTorchInitContainerTemplateFile, "pytorch-init-container-template-file",
 		config.PyTorchInitContainerTemplateFileDefault, "The template file for pytorch init container")
+	flag.IntVar(&config.Config.PyTorchInitContainerMaxTries, "pytorch-init-container-max-tries",
+		config.PyTorchInitContainerMaxTriesDefault, "The number of tries for the pytorch init container")
 
 	// MPI related flags
 	flag.StringVar(&config.Config.MPIKubectlDeliveryImage, "mpi-kubectl-delivery-image",
diff --git a/pkg/config/config.go b/pkg/config/config.go
index 2a8ed5e813..a8e3b4fbfc 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -19,6 +19,7 @@ var Config struct {
 	PyTorchInitContainerTemplateFile string
 	PyTorchInitContainerImage        string
 	MPIKubectlDeliveryImage          string
+	PyTorchInitContainerMaxTries     int
 }
 
 const (
@@ -28,6 +29,8 @@ const (
 	// PyTorchInitContainerTemplateFileDefault is the default template file for
 	// the pytorch init container.
 	PyTorchInitContainerTemplateFileDefault = "/etc/config/initContainer.yaml"
+	// PyTorchInitContainerMaxTriesDefault is the default number of tries for the pytorch init container.
+	PyTorchInitContainerMaxTriesDefault = 100
 	// MPIKubectlDeliveryImageDefault is the default image for launcher pod in MPIJob init container.
 	MPIKubectlDeliveryImageDefault = "mpioperator/kubectl-delivery:latest"
 )
diff --git a/pkg/controller.v1/pytorch/initcontainer.go b/pkg/controller.v1/pytorch/initcontainer.go
index 39dc7d4725..8c36e06a31 100644
--- a/pkg/controller.v1/pytorch/initcontainer.go
+++ b/pkg/controller.v1/pytorch/initcontainer.go
@@ -43,7 +43,7 @@ var (
     requests:
       cpu: 50m
       memory: 10Mi
-  command: ['sh', '-c', 'until nslookup {{.MasterAddr}}; do echo waiting for master; sleep 2; done;']`
+  command: ['sh', '-c', 'err=1;for i in $(seq {{.MaxTries}}); do if nslookup {{.MasterAddr}}; then err=0 && break; fi;echo waiting for master; sleep 2; done; exit $err']`
 	onceInitContainer sync.Once
 	icGenerator       *initContainerGenerator
 )
@@ -51,6 +51,7 @@ var (
 type initContainerGenerator struct {
 	template string
 	image    string
+	maxTries int
 }
 
 func getInitContainerGenerator() *initContainerGenerator {
@@ -58,6 +59,7 @@ func getInitContainerGenerator() *initContainerGenerator {
 		icGenerator = &initContainerGenerator{
 			template: getInitContainerTemplateOrDefault(config.Config.PyTorchInitContainerTemplateFile),
 			image:    config.Config.PyTorchInitContainerImage,
+			maxTries: config.Config.PyTorchInitContainerMaxTries,
 		}
 	})
 	return icGenerator
@@ -72,9 +74,11 @@ func (i *initContainerGenerator) GetInitContainer(masterAddr string) ([]corev1.C
 	if err := tpl.Execute(&buf, struct {
 		MasterAddr         string
 		InitContainerImage string
+		MaxTries           int
 	}{
 		MasterAddr:         masterAddr,
 		InitContainerImage: i.image,
+		MaxTries:           i.maxTries,
 	}); err != nil {
 		return nil, err
 	}


Test patch:
diff --git a/pkg/controller.v1/pytorch/initcontainer_test.go b/pkg/controller.v1/pytorch/initcontainer_test.go
index 3ead5855a9..e2f06cae1a 100644
--- a/pkg/controller.v1/pytorch/initcontainer_test.go
+++ b/pkg/controller.v1/pytorch/initcontainer_test.go
@@ -34,6 +34,7 @@ func TestInitContainer(t *testing.T) {
 
 	config.Config.PyTorchInitContainerImage = config.PyTorchInitContainerImageDefault
 	config.Config.PyTorchInitContainerTemplateFile = config.PyTorchInitContainerTemplateFileDefault
+	config.Config.PyTorchInitContainerMaxTries = config.PyTorchInitContainerMaxTriesDefault
 
 	testCases := []struct {
 		job         *kubeflowv1.PyTorchJob