axlearn icon indicating copy to clipboard operation
axlearn copied to clipboard

Alternative method of submitting jobs to DF Runner

Open remylouisew opened this issue 1 year ago • 2 comments

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.

remylouisew avatar Oct 17 '24 22:10 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?

remylouisew avatar Oct 18 '24 19:10 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?

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.

markblee avatar Oct 24 '24 22:10 markblee

Closing this PR due to inactivity. Please re-open or file a new PR if this is still important.

Ethanlm avatar Jul 23 '25 22:07 Ethanlm