add user_identity to task for elivated permissions

zooniverse · Dec 16, 2024 · f3f5039 · f3f5039
1 parent ab82fb9
commit f3f5039
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/bajor/batch/predictions.py b/bajor/batch/predictions.py
@@ -88,6 +88,12 @@ def create_batch_job(job_id, manifest_url, pool_id, checkpoint_target='ZOOBOT_CH
     job.job_preparation_task = batchmodels.JobPreparationTask(
         command_line=f'/bin/bash -c \"set -ex; {create_results_dir}; {copy_code_to_shared_dir}\"',
         constraints=batchmodels.TaskConstraints(max_task_retry_count=3),
+        user_identity = batchmodels.UserIdentity(
+           auto_user=batchmodels.AutoUserSpecification(
+              scope=batchmodels.AutoUserScope.task,
+              elevation_level=batchmodels.ElevationLevel.admin
+           )
+        ),
         #
         # A busted preparation task means the main task won't launch...ever!
         # and leave the node in a scaled state costing $$ ££

diff --git a/bajor/batch/train_finetuning.py b/bajor/batch/train_finetuning.py
@@ -116,6 +116,12 @@ def create_batch_job(job_id, manifest_container_path, pool_id, checkpoint_target
     job.job_preparation_task = batchmodels.JobPreparationTask(
         command_line=f'/bin/bash -c \"set -ex; {setup_pytorch_kernel_cache_dir}; {create_results_dir}; {copy_code_to_shared_dir}\"',
         constraints=batchmodels.TaskConstraints(max_task_retry_count=3),
+        user_identity = batchmodels.UserIdentity(
+           auto_user=batchmodels.AutoUserSpecification(
+              scope=batchmodels.AutoUserScope.task,
+              elevation_level=batchmodels.ElevationLevel.admin
+           )
+        ),
         #
         # A busted preparation task means the main task won't launch...ever!
         # and leave the node in a scaled state costing $$ ££