Fixing task ID replacement for MNP jobs on AWS Batch #2574
Conversation
| self._task_id.replace("control-", "") | ||
| + "-node-$AWS_BATCH_JOB_NODE_INDEX", | ||
| ) | ||
| # Fix: Only replace task ID in specific arguments, not in environment variables |
There was a problem hiding this comment.
thanks for the fix! @saikonen any suggestions on approaches that might allow us to not lean on string substitution (quite finicky)!
There was a problem hiding this comment.
Haven't touched this implementation in a bit so maybe there is a reason that the task-id patterns are kept in execute() but at first glance this seems (in the original implementation) like quite a late phase to be doing anything with the command.
Could a better place be when the cmd is not yet joined into a string, when we can operate on options separately? f.ex. batch_cli.py#251 and step_functions.py#925 ?
There was a problem hiding this comment.
I'm not 100% sure of the flow but I think that is possible if the flag --task-id is the only thing that needs to be changed. Looking at the command, I'm not sure if there is something else that requires us to modify the task ID (ex. MF_PATHSPEC). If task-id is the only thing, then I can definitely make that change.
There was a problem hiding this comment.
I've updated the code. I think it is difficult to fully rely on batch_cli.py#251 because we don't have access to $AWS_BATCH_JOB_NODE_INDEX as this would only be available at runtime for a given worker MNP node, I believe. But I've added more placeholders in batch_cli.py#251 that should make the regex more reliable. I'm less familiar with the step_functions file as I'm not using that right now.
Is it possible to get approval on this soon?
There was a problem hiding this comment.
the direction looks good as is so no further changes should be required. My final testing is still pending due to some unrelated infra issues at my end that I need to look into. Aiming to get this PR wrapped up by this week.
There was a problem hiding this comment.
were you able to have success with this PR on your end? I was finally able to test things and it seems that it is not working as expected. The parallel jobs do launch, but the main control node never finishes, instead timing out.
I would assume that the issue lies in https://github.com/Netflix/metaflow/blob/master/metaflow/plugins/aws/batch/batch_decorator.py#L412
where we poll the status of mapper tasks via the Metaflow client. This relies on task metadata, either through a metadata service, or in the absence of one, the data stored on local disk.
multi-node parallel jobs on batch are single tenant, so they do not share disk between each other. This would mean that a mapper task recording its finished_at locally (due to --metadata local) will not be accessible by the polling control task.
There was a problem hiding this comment.
Apologies @saikonen, I had pre-maturely terminated the jobs assuming it was ok. I've now run this end-to-end and confirmed that this works on my end! If we don't have a metadata provider, I've resorted to just using the S3 datastore to check MNP job statuses, which has worked smoothly for me.
…place with node-index
If we aren't using the Metaflow metadata service provider, Metaflow defaults to generating task IDs locally. But these task IDs are just simple integers based on how many tasks/steps there are and are sequentially incremented based on
new_task_idinmetaflow/plugins/metadata_providers/local.py. This presents a problem when we're doing AWS Batch MNP, since currently we try and mass replace based on the task ID in the secondary command. If this is a simple integer, this will replace many erroneous places.For example, if the task ID is "3", there could be many instances of "3" in the secondary command that then have many replacements with "-node-$AWS_BATCH_JOB_NODE_INDEX" when really we just want to replace the actual task ID.
Here, I've identified two places - the input task ID via
--task-idand the task ID inMF_PATHSPEC, that should be the only two places in the command that have the actual task ID in them that need replacing. It is better to have more specific regexes this way.Furthermore, if there is no metadata provider, I've added a new check for control MNP jobs to finish by checking the S3 datastore instead.