Add flag to pass the private image repository secret name (#523)

* fix issue #491

Signed-off-by: Piotr Pawłowski <[email protected]>

---------

Signed-off-by: Piotr Pawłowski <[email protected]>

Author: pawloch00

.github/workflows/reusable_workload_tests.yaml

@@ -30,6 +30,7 @@ on:
 
 env:
   WORKLOAD_NAME: xpktest-build-${{ github.run_attempt }}
+  PRIVATE_IMAGE_WORKLOAD_NAME: xpktest-build-priv-${{ github.run_attempt }}
   PATHWAYS_WORKLOAD_NAME: xpkpw-build-${{ github.run_attempt }}
 
 jobs:
@@ -83,8 +84,24 @@ jobs:
       run: python3 xpk.py workload list --cluster ${{inputs.cluster-name}} --zone=${{inputs.zone}} --wait-for-job-completion $PATHWAYS_WORKLOAD_NAME --timeout 300
     - name: List out the workloads on the cluster
       run: python3 xpk.py workload list --cluster ${{inputs.cluster-name}} --zone=${{inputs.zone}}
+    - name: Create a kubectl secret to use with private images
+      run: | 
+        kubectl create secret docker-registry gcr-key \
+        --docker-server=${{secrets.DOCKER_REPO_SERVER}}ubuntu2004 \
+        --docker-username=_json_key \
+        --docker-password='${{secrets.GCP_SA_KEY}}' \
+        --docker-email='${{secrets.GCP_SA_EMAIL}}'
+    - name: Run workload with private image
+      run: python xpk.py workload create --cluster ${{inputs.cluster-name}} --workload $PRIVATE_IMAGE_WORKLOAD_NAME  --command "echo foo"  --tpu-type=${{inputs.tpu-type}} --num-slices=1 --zone=${{inputs.zone}} --docker-image=${{secrets.DOCKER_REPO_SERVER}}ubuntu2004 --docker-image-pull-secret=gcr-key
+    - name: Wait for private image workload completion and confirm it succeeded
+      run: python3 xpk.py workload list --cluster ${{inputs.cluster-name}} --zone=${{inputs.zone}} --wait-for-job-completion $PRIVATE_IMAGE_WORKLOAD_NAME --timeout 300
+    - name: Delete kubectl secret
+      if: always()
+      run: kubectl delete secret gcr-key
     - name: Run xpk info
       run: python3 xpk.py info --cluster ${{inputs.cluster-name}} --zone=${{inputs.zone}} | tee output.txt | grep -P "^(?=.*QUEUE)(?=.*PENDING_WORKLOADS)(?=.*ADMITTED_WORKLOADS)(?=.*1x${{inputs.tpu-type}}:google.com/tpu)(?=.*cpu-user:cpu)(?=.*cpu-user:memory)" || (echo 'Invalid command output' && cat output.txt && exit 1)
+    - name: Delete the private image workload on the cluster
+      run: python3 xpk.py workload delete --workload $PRIVATE_IMAGE_WORKLOAD_NAME --cluster ${{inputs.cluster-name}} --zone=${{inputs.zone}}
     - name: Delete the workload on the cluster
       run: python3 xpk.py workload delete --workload $WORKLOAD_NAME --cluster ${{inputs.cluster-name}} --zone=${{inputs.zone}}
     - name: Delete the Pathways workload on the cluster

src/xpk/commands/workload.py

@@ -125,6 +125,8 @@
                 {storage_annotations}
             spec:
               schedulerName: {args.scheduler}
+              imagePullSecrets:
+              - name: {args.docker_image_pull_secret}
               restartPolicy: Never
               {affinity}
               nodeSelector:
@@ -177,6 +179,8 @@
               {gpu_scheduler}
               priorityClassName: {args.priority}
               restartPolicy: Never
+              imagePullSecrets:
+              - name: {args.docker_image_pull_secret}
               hostNetwork: true
               dnsPolicy: ClusterFirstWithHostNet
               terminationGracePeriodSeconds: {args.termination_grace_period_seconds}
@@ -220,6 +224,8 @@
             spec:
               priorityClassName: {args.priority}
               restartPolicy: Never
+              imagePullSecrets:
+              - name: {args.docker_image_pull_secret}
               dnsPolicy: ClusterFirstWithHostNet
               terminationGracePeriodSeconds: {args.termination_grace_period_seconds}
               serviceAccountName: {service_account}

src/xpk/parser/workload.py

@@ -728,6 +728,14 @@ def add_shared_workload_docker_image_arguments(args_parsers):
             ' directly by the xpk workload.'
         ),
     )
+    custom_parser.add_argument(
+        '--docker-image-pull-secret',
+        type=str,
+        help=(
+            'Name of the secret that will be used to pull image from'
+            ' private repository'
+        ),
+    )
 
 
 def add_shared_workload_create_tensorboard_arguments(args_parsers):