refactor to use the e2e test script

[michelle-yooh]

Description

This PR refactors the maxtext_moe_gpu_e2e tests to use the MaxText e2e test script instead of the direct command. It requires AI-Hypercomputer/maxtext#1191 to be merged first.

Tests

Please describe the tests that you ran on Cloud VM to verify changes.

Instruction and/or command lines to reproduce your tests: ...

List links for your tests (use go/shortn-gen for any internal link): ...

Screenshot of the test result
http://shortn/_dNMK6pqBn9

bugs for the failed tests
pinned test on 2 A3 nodes: http://shortn/_WN4xgkG3Hj
pinned test on 2 A3+ nodes: http://shortn/_zKqVme3CFO

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run one-shot tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[michelle-yooh]

@RissyRan @parambole I've moved the functions to xlml/utils/gpu.py and verified that the resizing works as expected. http://shortn/_ttKViB2Hv7 (Please disregard the test failures as it used the latest MaxText image which doesn't have the test script yet)