Skip to content

add aws support to ai-inference in pkb #6201

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kiryl-filatau wants to merge 4 commits into
base: master
Choose a base branch
from
Open

add aws support to ai-inference in pkb #6201

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
7ed5fef
add aws support to ai-inference in pkb with temporary link to wg-serv…
kiryl-filatau Nov 5, 2025
013cb43
Merge branch 'master' into feature/add-aws-support-to-ai-inference-pkb
kiryl-filatau Nov 18, 2025
b2c0067
added AWS-only condition for gpu nodepool manifest, moved wg-serving …
kiryl-filatau Nov 19, 2025
c2f11fc
move WG Serving repo flags to wg_serving_inference_server and use fla…
kiryl-filatau Nov 24, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 27 additions & 0 deletions perfkitbenchmarker/data/container/kubernetes_ai_inference/aws-gpu-nodepool.yaml.j2
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
apiVersion: karpenter.sh/v1
kind: NodePool
metadata:
name: {{ gpu_nodepool_name | default('gpu') }}
spec:
disruption:
consolidateAfter: {{ gpu_consolidate_after | default('1h') }}
consolidationPolicy: {{ gpu_consolidation_policy | default('WhenEmpty') }}
template:
spec:
nodeClassRef:
group: eks.amazonaws.com
kind: NodeClass
name: {{ karpenter_nodeclass_name | default('default') }}
requirements:
- key: karpenter.sh/capacity-type
operator: In
values: {{ gpu_capacity_types | default(['on-demand']) }}
- key: kubernetes.io/arch
operator: In
values: {{ gpu_arch | default(['amd64']) }}
- key: eks.amazonaws.com/instance-family
operator: In
values: {{ gpu_instance_families | default(['g6','g6e']) }}
taints:
- key: {{ gpu_taint_key | default('nvidia.com/gpu') }}
effect: NoSchedule
4 changes: 4 additions & 0 deletions perfkitbenchmarker/data/container/kubernetes_ai_inference/serving_catalog_cli.yaml.j2
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,10 @@ spec:
value: {{generate_args}}
- name: OUTPUT_FILE
value: /output/output.yaml
- name: WG_SERVING_REPO_URL
value: "{{ wg_serving_repo_url }}"
- name: WG_SERVING_REPO_BRANCH
value: "{{ wg_serving_repo_branch }}"
volumeMounts:
- name: output
mountPath: /output
Expand Down
12 changes: 12 additions & 0 deletions perfkitbenchmarker/linux_benchmarks/kubernetes_ai_inference_benchmark.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,18 @@
' --k8s_ai_inference_extra_lpg_args=KEY_OF_ARG=VALUE_OF_ARG',
)

flags.DEFINE_string(
'wg_serving_repo_url',
'https://github.com/kubernetes-sigs/wg-serving',
'Repository URL for WG-Serving used by serving_catalog_cli.',
)

flags.DEFINE_string(
'wg_serving_repo_branch',
'main',
Copy link
Collaborator

@hubatish hubatish Nov 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice. I was wondering if you'd have to leave off values altogether with like a jinja if, but this totally makes sense as defaults.

'Git branch of the WG-Serving repository.',
)

lpg_extra_args = {}


Expand Down
15 changes: 15 additions & 0 deletions perfkitbenchmarker/resources/kubernetes/wg_serving_inference_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -569,6 +569,19 @@ def _InjectDefaultHuggingfaceToken(self) -> None:

def _GetInferenceServerManifest(self) -> str:
"""Generates and retrieves the inference server manifest content."""
# Ensure GPU capacity exists before scheduling GPU workloads
Copy link
Collaborator

@hubatish hubatish Jan 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On my end this looks merged in:
https://screenshot.googleplex.com/AREUuCaL76yAdQg

Ooh ok actually perhaps this has made it in but been superceded by follow up PRs? https://screenshot.googleplex.com/BNVuvNLwNKckBF5 is the current state of wg_serving_inference_server.py & it does have a cloud == 'AWS' section.
I suppose we simply close this PR?

if FLAGS.cloud == 'AWS':
self.cluster.ApplyManifest(
'container/kubernetes_ai_inference/aws-gpu-nodepool.yaml.j2',
gpu_nodepool_name='gpu',
gpu_consolidate_after='1h',
gpu_consolidation_policy='WhenEmpty',
karpenter_nodeclass_name='default', # must exist already
gpu_capacity_types=['on-demand'],
gpu_arch=['amd64'],
gpu_instance_families=['g6','g6e'],
gpu_taint_key='nvidia.com/gpu',
)
generate_args = {
'kind': 'core/deployment',
'model-server': self.spec.model_server,
Expand All @@ -581,6 +594,8 @@ def _GetInferenceServerManifest(self) -> str:
self.cluster.ApplyManifest(
'container/kubernetes_ai_inference/serving_catalog_cli.yaml.j2',
image_repo=FLAG_IMAGE_REPO.value,
wg_serving_repo_url=FLAGS.wg_serving_repo_url,
Copy link
Collaborator

@hubatish hubatish Nov 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok, flag location is always tricky though. Let's put the flags in this file (wg_serving_inference_server) & use flagholders with eg REPO_URL = FLAGS.DEFINE.... REPO_URL.value.

Copy link
Collaborator Author

@kiryl-filatau kiryl-filatau Nov 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

changed

wg_serving_repo_branch=FLAGS.wg_serving_repo_branch,
generate_args=' '.join(
[f'--{k} {v}' for k, v in generate_args.items()]
),
Expand Down