ai-ml/gke-a100-jax/kubernetes/kustomization.yaml

# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


namespace: default 

resources:
- job.yaml
- service.yaml

apiVersion: kustomize.config.k8s.io/v1beta1
kind: Kustomization
patches:
- path: job-patch.yaml
- patch: |-
    - op: replace
      path: /spec/completions
      value: 32
    - op: replace
      path: /metadata/name
      value: jax-hello-world
    - op: add 
      path: /spec/template/spec/nodeSelector
      value: 
         cloud.google.com/gke-accelerator: 
           nvidia-a100-80gb
  target:
    kind: Job
    version: v1

images:
- name: gcr.io/<<PROJECT>>/jax/hello:latest
  newName: gcr.io/<<PROJECT>>/jax/hello
  newTag: latest

replacements:
- source:
    kind: Job
    version: v1
  targets:
  - select:
      kind: Service
      version: v1
    fieldPaths:
    - spec.selector.job-name
  - select:
      kind: Job
      version: v1
    fieldPaths:
    - spec.template.spec.containers.[name=jax-worker].args.3 
- source:
    kind: Job
    version: v1
    fieldPath: spec.template.spec.subdomain
  targets:
  - select:
      kind: Service
      version: v1
    fieldPaths:
    - metadata.name
  - select:
      kind: Job
      version: v1
    fieldPaths:
    - spec.template.spec.containers.[name=jax-worker].args.5 
- source:
    kind: Job
    version: v1
    fieldPath: spec.completions
  targets:
  - select: 
      kind: Job
      version: v1
    fieldPaths:
    - spec.template.spec.containers.[name=jax-worker].args.1
    - spec.parallelism
- source:
    kind: Job
    version: v1
    fieldPath: spec.template.spec.containers.[name=jax-worker].ports.0.containerPort
  targets:
  - select: 
      kind: Job
      version: v1
    fieldPaths:
    - spec.template.spec.containers.[name=jax-worker].args.7