add benchmark to the new file instruction grain metadata impl

[sadikneipp]

No description provided.

[gemini-code-assist]

Code Review

This pull request introduces a new standalone JobSet configuration (benchmarks/jobset_benchmark_file_instruction.yaml) to run the MaxText FileInstruction initialization benchmark on Cloud TPU v5e-32. The review feedback highlights several issues to address: a syntax error in the GRPC_SERVER_ADDRESS environment variable due to triple single quotes, an unreliable _sigterm trap handler in the coordinator container, the use of fragile personal/dev container images, hardcoded Google-internal GCS buckets that will cause permission issues for external users, and a fragile hardcoded sleep timer for worker pod initialization.

[gemini-code-assist]

The environment variable GRPC_SERVER_ADDRESS has triple single quotes ('''0.0.0.0:50051'''). In YAML, this will be parsed as a string with literal single quotes around the IP address ('0.0.0.0:50051'), which will cause gRPC server binding or connection to fail. It should be defined without the extra single quotes.

                value: "0.0.0.0:50051"

[gemini-code-assist]

The _sigterm trap handler uses $! directly and runs in a subshell () without exiting. If a SIGTERM is received, the coordinator shell will run the trap but continue executing subsequent commands instead of propagating the termination. Additionally, using $! directly can be unreliable if other background processes are started. It is safer and more robust to use the captured PID variable and explicitly exit, matching the clean design of the worker's _term trap.

                _sigterm() {
                  echo "Caught SIGTERM signal! Forwarding to child..."
                  if [ -n "${PID:-}" ]; then
                    kill -SIGTERM "$PID" 2>/dev/null
                  fi
                  exit 1
                }
                trap _sigterm SIGTERM;

[gemini-code-assist]

The container image points to a personal/dev registry path (ksadi/maxtext-runner). Dev images can be deleted or modified at any time, which makes the benchmark fragile and prone to breaking. Please use an official or shared release image path instead.

[gemini-code-assist]

The container image points to a personal/dev registry path (ksadi/maxtext-colocated-python). Dev images can be deleted or modified at any time, which makes the benchmark fragile and prone to breaking. Please use an official or shared release image path instead.

[gemini-code-assist]

The benchmark configuration uses hardcoded Google-internal GCS buckets (gs://tess-pin-checkpointing-us-central1/... and gs://tess-pin-dataloading-us-central1/...). External users running this benchmark will encounter permission denied errors. Consider parameterizing these paths or adding a comment/documentation on how users can configure their own GCS buckets.

[gemini-code-assist]

Using a hardcoded sleep 150 to wait for worker pods to initialize is fragile. If the worker pods take longer to initialize (e.g., due to image pulling or scheduling delays), the coordinator will start the benchmark prematurely and fail. Consider implementing a dynamic check to poll the readiness of the worker pods or endpoints before starting the benchmark.