Skip to content

Commit

Permalink
Add exports property to the config to add env vars to export to the job
Browse files Browse the repository at this point in the history
  • Loading branch information
BenGalewsky committed Dec 7, 2024
1 parent 28d1478 commit 49389c6
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions mlflow_slurm/templates/sbatch_template.sh
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
{% if config.account %}
#SBATCH --account={{ config.account }}
{% endif %}
#SBATCH --export=MLFLOW_TRACKING_URI,MLFLOW_TRACKING_TOKEN,MLFLOW_TRACKING_USERNAME,MLFLOW_TRACKING_PASSWORD,MLFLOW_S3_ENDPOINT_URL,AWS_SECRET_ACCESS_KEY,AWS_ACCESS_KEY_ID
#SBATCH --export=MLFLOW_TRACKING_URI,MLFLOW_TRACKING_TOKEN,MLFLOW_TRACKING_USERNAME,MLFLOW_TRACKING_PASSWORD,MLFLOW_S3_ENDPOINT_URL,AWS_SECRET_ACCESS_KEY,AWS_ACCESS_KEY_ID,{{ config.exports |join(',') }}
{% if config.gpus_per_node %}
#SBATCH --gpus-per-node={{ config.gpus_per_node }}
{% endif %}
Expand Down Expand Up @@ -43,7 +43,7 @@ export {{ env }}
export MLFLOW_RUN_ID={{ run_id }}
echo "job is starting on `hostname`"
{% if config.nodes %}
srun --export=ALL /bin/bash -c '{{ command }}'
srun --export=ALL /bin/bash -c '{{ command }}'
{% else %}
{{ command }}
{% endif %}

0 comments on commit 49389c6

Please sign in to comment.