Is it possible to use @metaflow_ray with foreach on AWS Batch?

[vymao]

I noticed that the multinode case for AWS Batch (which I believe is what is called for @metaflow_ray) seems to rely on UBF_CONTROL , which makes me think that we can't use call a multinode batch step in a foreach task (because then presumably we would be using UBF_TASK). Is this correct?

I'm wondering because my use case has multiple job groups that I would like to submit jobs for, where each job group is a parallelizable set of subjobs. I wanted distinction between the individual job groups themselves so I thought combining the foreach and the metaflow_ray decorator would be useful

[savingoyal]

yes, @metaflow-ray does work on AWS Batch using multi-node. you can nest metaflow-ray within a foreach.

[vymao]

Got it. Do you know what might be causing this error then? It seems to be referencing a path called s3://valthos-compute//jobs/DesignFlow/data/65/654-node-$AWS_BATCH_JOB_NODE_INDEXfc01a321fbdc7665d25d54-node-$AWS_BATCH_JOB_NODE_INDEXd654-node-$AWS_BATCH_JOB_NODE_INDEXe7164-node-$AWS_BATCH_JOB_NODE_INDEX0a4-node-$AWS_BATCH_JOB_NODE_INDEX912 but the only path that exists is s3://valthos-compute/jobs/DesignFlow/data/65/654fc01a321fbdc7665d25d54d654e71640a4912, which itself is empty.

Downloading code package...
--
Traceback (most recent call last):
File "<string>", line 1, in <module>
File "/usr/local/lib/python3.11/site-packages/botocore/context.py", line 123, in wrapper
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/boto3/s3/inject.py", line 223, in download_file
return transfer.download_file(
^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/boto3/s3/transfer.py", line 406, in download_file
future.result()
File "/usr/local/lib/python3.11/site-packages/s3transfer/futures.py", line 111, in result
return self._coordinator.result()
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/s3transfer/futures.py", line 287, in result
raise self._exception
File "/usr/local/lib/python3.11/site-packages/s3transfer/tasks.py", line 272, in _main
self._submit(transfer_future=transfer_future, **kwargs)
File "/usr/local/lib/python3.11/site-packages/s3transfer/download.py", line 355, in _submit
response = client.head_object(
^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/botocore/client.py", line 601, in _api_call
return self._make_api_call(operation_name, kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/botocore/context.py", line 123, in wrapper
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/botocore/client.py", line 1074, in _make_api_call
raise error_class(parsed_response, operation_name)
botocore.exceptions.ClientError: An error occurred (404) when calling the HeadObject operation: Not Found
Failed to download code package from s3://valthos-compute/jobs/DesignFlow/data/65/654-node-$AWS_BATCH_JOB_NODE_INDEXfc01a321fbdc7665d25d54-node-$AWS_BATCH_JOB_NODE_INDEXd654-node-$AWS_BATCH_JOB_NODE_INDEXe7164-node-$AWS_BATCH_JOB_NODE_INDEX0a4-node-$AWS_BATCH_JOB_NODE_INDEX912 after 6 tries. Exiting...

[savingoyal]

can you share your a reproducible example? the path it is trying to read from is incorrect.

[vymao]

Sure, this is the portion of my code that creates these jobs.

   @step
    def prepare_design_inputs(self):
        # Load base defaults
        base_filters = read_defaults(self.default_filter_file)
        print("Done reading base filters.")
        base_advanced_settings = read_defaults(self.default_advanced_settings_file)
        print("Done reading base advanced settings.")
        base_target_settings = read_defaults(self.default_target_settings_file)
        print("Done reading base target settings.")

        # First expand job combinations
        expanded_config = expand_all_job_combinations(self.job_config)

        print(expanded_config)

        # Create individual overrides for each job combination
        self.job_overrides = {}

        # Handle the new nested structure: original_job_name -> {expanded_job_name: job_config}
        for original_job_name, job_combinations in expanded_config["jobs"].items():
            self.job_overrides[original_job_name] = {}
            for expanded_job_name, job_config in job_combinations.items():
                # Use helper functions to validate and override settings

                job_advanced_settings = override_advanced_settings(
                    base_advanced_settings, job_config, expanded_job_name
                )
                print("Done overriding advanced settings: ", job_advanced_settings)
                job_target_settings = override_target_settings(
                    base_target_settings, job_config, expanded_job_name
                )
                print("Done overriding target settings: ", job_target_settings)

                # Store the overrides for this specific job combination
                self.job_overrides[original_job_name][expanded_job_name] = {
                    "filters": base_filters.copy(),  # Don't override the base filters for now.
                    "advanced_settings": job_advanced_settings,
                    "target_settings": job_target_settings,
                    "original_job_name": original_job_name,  # Keep reference to original job
                }

        # Prepare job data for foreach - each original job name with its expanded jobs
        self.job_groups = []
        for original_job_name, job_combinations in self.job_overrides.items():
            self.job_groups.append(
                {
                    "original_job_name": original_job_name,
                    "expanded_jobs": job_combinations,
                    "num_jobs": len(job_combinations),
                }
            )

        if not hasattr(self, "job_groups"):
            raise RuntimeError("job_groups not computed; earlier code failed")

        self.next(self.design_job_group, foreach="job_groups")

    @step
    def design_job_group(self):
        """Prepare for Ray execution by determining cluster size."""
        job_group = self.input
        original_job_name = job_group["original_job_name"]
        expanded_jobs = job_group["expanded_jobs"]
        num_jobs = job_group["num_jobs"]

        # Store job group data for the Ray step
        self.original_job_name = original_job_name
        self.expanded_jobs = expanded_jobs
        self.num_jobs = num_jobs

        print(f"Preparing Ray cluster for job group: {original_job_name}")
        print(f"Number of expanded jobs in this group: {num_jobs}")

        # Pass the number of jobs as num_parallel to the Ray step. We add 1 to account for the
        # head node.
        self.next(self.execute_ray_jobs, num_parallel=num_jobs)

    # @secrets(sources=[EnvConfig.get("github_secret_arn")])
    @batch(
        image=ImageConfig.get_image_url("gpu-base"),
        queue=InstanceConfig.get_queue_name("L40S"),
        gpu=1,
    )
    @metaflow_ray(
        # 20 mins timeout for all workers to join
        all_nodes_started_timeout=1200
    )
    @environment(
        vars={
            "GIT_BRANCH": get_current_branch(),
            "GIT_REPO": VALTHOS_MAIN_REPO,
        }
    )
    @step
    def execute_ray_jobs(self):
        """Execute jobs using Ray with dynamically sized cluster."""
        original_job_name = self.original_job_name
        expanded_jobs = self.expanded_jobs
        num_jobs = self.num_jobs

        print(f"Starting Ray cluster for job group: {original_job_name}")
        print(f"Number of expanded jobs in this group: {num_jobs}")

        # Get the specific job for this parallel node
        expanded_job_names = list(expanded_jobs.keys())
        node_index = current.parallel.node_index

        future_results = []
        self.job_results = []

[vymao]

Digging myself, I can only find reference to the code package path in question here, where $AWS_BATCH_JOB_NODE_INDEX is used.

[vymao]

I believe there is a mismatch between what is passed into the job through METAFLOW_CODE_URL and what the actual secondary command is. One of my nodes has a command like

["bash","-c","true && mkdir -p $PWD/.logs && export PYTHONUNBUFFERED=x MF_PATHSPEC=DesignFlow/1756138917013236/execute_ray_jobs/4 MF_DATASTORE=s3 MF_ATTEMPT=0 MFLOG_STDOUT=$PWD/.logs/mflog_stdout MFLOG_STDERR=$PWD/.logs/mflog_stderr && mflog(){ T=$(date -u -Ins|tr , .); echo \"[MFLOG|0|${T:0:26}Z|task|$T]$1\" >> $MFLOG_STDOUT; echo $1;  } && flush_mflogs(){ python -m metaflow.mflog.save_logs; } && mflog 'Setting up task environment.' && if [ -z $METAFLOW_SKIP_INSTALL_DEPENDENCIES ]; then python -m pip install -qqq --no-compile --no-cache-dir --disable-pip-version-check boto3 requests; fi && mkdir metaflow && cd metaflow && mkdir .metaflow && i=0; while [ $i -le 5 ]; do mflog 'Downloading code package...'; python -c 'import boto3, os; ep=os.getenv(\"METAFLOW_S3_ENDPOINT_URL\"); boto3.client(\"s3\", **({\"endpoint_url\":ep} if ep else {})).download_file(\"valthos-compute\", \"jobs/DesignFlow/data/65/654fc01a321fbdc7665d25d54d654e71640a4912\", \"job.tar\")' && mflog 'Code package downloaded.' && break; sleep 10; i=$((i+1)); done && if [ $i -gt 5 ]; then mflog 'Failed to download code package from s3://valthos-compute/jobs/DesignFlow/data/65/654fc01a321fbdc7665d25d54d654e71640a4912 after 6 tries. Exiting...' && exit 1; fi && TAR_OPTIONS='--warning=no-timestamp' tar -xzf job.tar -C . && export PYTHONPATH=$(pwd)/.mf_code:$(printenv PYTHONPATH) && export METAFLOW_EXTRACTED_ROOT=$(pwd) && mflog 'Task is starting.' && flush_mflogs && (python -u design.py --quiet --datastore s3 --metadata local --environment local --event-logger nullSidecarLogger --monitor nullSidecarMonitor --datastore-root s3://valthos-compute/jobs --package-suffixes .py,.R,.RDS --pylint step execute_ray_jobs --run-id 1756138917013236 --task-id 4 --input-paths ${METAFLOW_INPUT_PATHS_0} --retry-count 0 --max-user-code-retries 0 --namespace user:victormaomain --ubf-context ubf_control --split-index 0 --split-index 0) 1>> >(python -m metaflow.mflog.tee task $MFLOG_STDOUT) 2>> >(python -m metaflow.mflog.tee task $MFLOG_STDERR >&2); c=$?; python -m metaflow.mflog.save_logs; exit $c"]

["bash","-c","true && mkdir -p $PWD/.logs && export PYTHONUNBUFFERED=x MF_PATHSPEC=DesignFlow/1756138917013236/execute_ray_jobs/4-node-$AWS_BATCH_JOB_NODE_INDEX MF_DATASTORE=s3 MF_ATTEMPT=0 MFLOG_STDOUT=$PWD/.logs/mflog_stdout MFLOG_STDERR=$PWD/.logs/mflog_stderr && mflog(){ T=$(date -u -Ins|tr , .); echo \"[MFLOG|0|${T:0:26}Z|task|$T]$1\" >> $MFLOG_STDOUT; echo $1;  } && flush_mflogs(){ python -m metaflow.mflog.save_logs; } && mflog 'Setting up task environment.' && if [ -z $METAFLOW_SKIP_INSTALL_DEPENDENCIES ]; then python -m pip install -qqq --no-compile --no-cache-dir --disable-pip-version-check boto3 requests; fi && mkdir metaflow && cd metaflow && mkdir .metaflow && i=0; while [ $i -le 5 ]; do mflog 'Downloading code package...'; python -c 'import boto3, os; ep=os.getenv(\"METAFLOW_S3_ENDPOINT_URL\"); boto3.client(\"s3\", **({\"endpoint_url\":ep} if ep else {})).download_file(\"valthos-compute\", \"jobs/DesignFlow/data/65/654-node-$AWS_BATCH_JOB_NODE_INDEXfc01a321fbdc7665d25d54-node-$AWS_BATCH_JOB_NODE_INDEXd654-node-$AWS_BATCH_JOB_NODE_INDEXe7164-node-$AWS_BATCH_JOB_NODE_INDEX0a4-node-$AWS_BATCH_JOB_NODE_INDEX912\", \"job.tar\")' && mflog 'Code package downloaded.' && break; sleep 10; i=$((i+1)); done && if [ $i -gt 5 ]; then mflog 'Failed to download code package from s3://valthos-compute/jobs/DesignFlow/data/65/654-node-$AWS_BATCH_JOB_NODE_INDEXfc01a321fbdc7665d25d54-node-$AWS_BATCH_JOB_NODE_INDEXd654-node-$AWS_BATCH_JOB_NODE_INDEXe7164-node-$AWS_BATCH_JOB_NODE_INDEX0a4-node-$AWS_BATCH_JOB_NODE_INDEX912 after 6 tries. Exiting...' && exit 1; fi && TAR_OPTIONS='--warning=no-timestamp' tar -xzf job.tar -C . && export PYTHONPATH=$(pwd)/.mf_code:$(printenv PYTHONPATH) && export METAFLOW_EXTRACTED_ROOT=$(pwd) && mflog 'Task is starting.' && flush_mflogs && (python -u design.py --quiet --datastore s3 --metadata local --environment local --event-logger nullSidecarLogger --monitor nullSidecarMonitor --datastore-root s3://valthos-compute/obs --package-suffixes .py,.R,.RDS --pylint step execute_ray_jobs --run-id 1756138917013236 --task-id 4-node-$AWS_BATCH_JOB_NODE_INDEX --input-paths ${METAFLOW_INPUT_PATHS_0} --retry-count 0 --max-user-code-retries 0 --namespace user:victormaomain --ubf-context ubf_task --split-index 0 --split-index $AWS_BATCH_JOB_NODE_INDEX) 1>> >(python -m metaflow.mflog.tee task $MFLOG_STDOUT) 2>> >(python -m metaflow.mflog.tee task $MFLOG_STDERR >&2); c=$?; python -m metaflow.mflog.save_logs; exit $c"]

However, both nodes have the same METAFLOW_CODE_URL as s3://valthos-compute/jobs/DesignFlow/data/65/654fc01a321fbdc7665d25d54d654e71640a4912

[vymao]

I believe I found the error: #2574