Alternative method of submitting jobs to DF Runner

[remylouisew]

Running a dataflow job via the axlearn gcp vm start command is not necessary or intuative for non-apple users. Additionally, if you are already running your commands from a VM (e.g from a remote desktop), this process does not work. I recognize that there are still scenarios where you would want to launch your dataflow jobs from a VM, so rather than replacing that ability, I am adding an alternative.

In order to submit jobs to the Dataflow runner without using ‘axlearn gcp vm start’, changes to the quoting behavior of dataflow.py are necessary. Unfortunately, there’s not an obviously elegant way to provide two versions of dataflow.py, so if you can think of a better option, please let me know.

What I’ve done is this: dataflow.py remains as it was, and I’m adding dataflow.alt.py. In the directions, I’ve added instructions to replace the original module if the user wants to submit jobs to the dataflow runner without needing to create a VM.

Additional note: PR #711 makes a change to the quoting behavior that allows the submission of dataflow jobs without ‘axlearn gcp vm start’, however this fix will not work for any commands that include parameters that require quotes, e.g --dataflow_service_options='worker_accelerator=type:nvidia-tesla-t4;count:1;install-nvidia-driver'
The dataflow.alt.py module DOES work these kinds of parameters.

[markblee]

Thanks @remylouisew -- IIUC, the main issue is with flag escaping. Aside from finding a more generic fix for flag parsing, maybe we can add support for loading from flagfiles, which would avoid the duplicate code and manual renaming step. (It seems that flag escaping is painful for users anyway.) WDYT?

[remylouisew]

@markblee The flag escaping has indeed been painful! I am not familiar with the process of loading from flagfiles, could you elaborate on how it would be implemented?

[markblee]

@markblee The flag escaping has indeed been painful! I am not familiar with the process of loading from flagfiles, could you elaborate on how it would be implemented?

absl has builtin support for flagfiles: https://abseil.io/docs/python/guides/flags#a-note-about---flagfile
So either we can accept a flagfile directly at the top-level axlearn gcp dataflow command, which is read and then flags forwarded to the user command; or the user script can use flagfiles directly. Let me know if additional clarifications are helpful.