Support enable_continuous_profiling option in the Tunix JAX profiler wrapper

[youchunni]

Description

Support enable_continuous_profiling option in the Tunix JAX profiler wrapper (sft/profiler.py).

When enable_continuous_profiling=True is configured in profiler options, this change sets advanced_configuration["enable_continuous_profiling"] = True in jax.profiler.ProfileOptions(). This configures JAX profiler to stream smaller trace chunks into .xplane.riegeli record containers, avoiding the 2GB Protocol Buffer size limit for large profiles.

Tested with the MaxText RL continuous profiling workflow. Related MaxText PR: AI-Hypercomputer/maxtext#5061

Note:

• This flag can be configured even when set_profile_options=False. MaxText sets set_profile_options=False to disable custom tracer level overrides, while enable_continuous_profiling is configured in profiler_options.
• This flag is only effective with JAX version >=0.11.1.

Checklist

• I have verified that my change does not break existing code and all unit tests pass.
• I have added all appropriate doc-strings/documentation.
• My PR is based on the latest changes of the main branch (if unsure, rebase the code).
• I have signed the Contributor License Agreement.
• I have followed Contribution Guidelines.