diff --git a/tests/common/environment.go b/tests/common/environment.go index eaafab52c..73654306f 100644 --- a/tests/common/environment.go +++ b/tests/common/environment.go @@ -69,6 +69,15 @@ func GetNotebookUserName(t Test) string { return name } +// GetNotebookUserNameFromEnv returns NOTEBOOK_USER_NAME when provided. +// Otherwise it falls back to resolving the user identity from the provided token. +func GetNotebookUserNameFromEnv(t Test, token string) string { + if name, ok := os.LookupEnv(notebookUserName); ok && strings.TrimSpace(name) != "" { + return name + } + return GenerateNotebookUserNameFromToken(t, token) +} + func GetNotebookUserToken(t Test) string { token, ok := os.LookupEnv(notebookUserToken) if !ok { @@ -77,6 +86,15 @@ func GetNotebookUserToken(t Test) string { return token } +// GetNotebookUserTokenFromEnv returns NOTEBOOK_USER_TOKEN when provided. +// Otherwise it falls back to generating a token from username/password. +func GetNotebookUserTokenFromEnv(t Test) string { + if token, ok := os.LookupEnv(notebookUserToken); ok && strings.TrimSpace(token) != "" { + return token + } + return GenerateNotebookUserToken(t) +} + func GetNotebookUserPassword(t Test) string { password, ok := os.LookupEnv(notebookUserPassword) if !ok { @@ -85,7 +103,7 @@ func GetNotebookUserPassword(t Test) string { return password } -// GenerateNotebookUserToken generates an OpenShift token using oc login with username and password +// GenerateNotebookUserToken generates an OpenShift token using oc login with username and password. func GenerateNotebookUserToken(t Test) string { userName := GetNotebookUserName(t) password := GetNotebookUserPassword(t) @@ -128,12 +146,39 @@ func GenerateNotebookUserToken(t Test) string { return strings.TrimSpace(string(out)) } +// GenerateNotebookUserNameFromToken resolves the username bound to a bearer token. +func GenerateNotebookUserNameFromToken(t Test, token string) string { + if strings.TrimSpace(token) == "" { + t.T().Fatalf("Cannot resolve Notebook username from token: token is empty") + } + + cmd := exec.Command( + "oc", "whoami", + "--token="+token, + "--server="+GetOpenShiftApiUrl(t), + "--insecure-skip-tls-verify=true", + ) + out, err := cmd.Output() + if err != nil { + if exitError, ok := err.(*exec.ExitError); ok { + t.T().Logf("Error running 'oc whoami' command: %v\n", exitError) + t.T().Logf("Output: %s\n", out) + t.T().Logf("Error output: %s\n", exitError.Stderr) + } else { + t.T().Logf("Error running 'oc whoami' command: %v\n", err) + } + t.T().FailNow() + } + + return strings.TrimSpace(string(out)) +} + func GetNotebookImage(t Test) string { - notebook_image, ok := os.LookupEnv(notebookImage) - if !ok { + notebookImageValue, ok := os.LookupEnv(notebookImage) + if !ok || strings.TrimSpace(notebookImageValue) == "" { t.T().Fatalf("Expected environment variable %s not found, please use this environment variable to specify image of the Notebook.", notebookImage) } - return notebook_image + return notebookImageValue } func GetTestTier(t Test) (string, bool) { diff --git a/tests/common/support/kueue_operator.go b/tests/common/support/kueue_operator.go index 046db1564..277a8af7f 100644 --- a/tests/common/support/kueue_operator.go +++ b/tests/common/support/kueue_operator.go @@ -17,6 +17,9 @@ limitations under the License. package support import ( + "encoding/json" + "fmt" + "strings" "time" "github.com/onsi/gomega" @@ -132,9 +135,72 @@ func VerifyKueueReady(test Test, expectedFrameworks ...string) { for _, framework := range expectedFrameworks { test.T().Logf("Verifying %s framework is present in Kueue CR...", framework) - test.Eventually(KueueCR(test, KueueCRName), TestTimeoutShort).Should( - gomega.WithTransform(KueueCRFrameworks, gomega.ContainElement(framework)), + // Capture baseline diagnostics before framework check polling starts. + dumpKueueDiagnostics(test, "before-framework-check", framework) + test.Eventually(func(g gomega.Gomega) bool { + kueue, err := GetKueueCR(test, KueueCRName) + g.Expect(err).NotTo(gomega.HaveOccurred()) + frameworks := KueueCRFrameworks(kueue) + if !containsString(frameworks, framework) { + // Update snapshot while polling so failures leave a concrete artifact. + dumpKueueDiagnostics(test, "framework-missing", framework) + return false + } + return true + }, TestTimeoutShort).Should( + gomega.BeTrue(), + "Expected framework '%s' to be present in Kueue CR", + framework, ) test.T().Logf("%s framework is present in Kueue CR", framework) } } + +func dumpKueueDiagnostics(test Test, stage string, expectedFramework string) { + test.T().Helper() + + kueue, err := GetKueueCR(test, KueueCRName) + if err != nil { + WriteToOutputDir( + test, + fmt.Sprintf("kueue-diagnostics-%s", sanitizeFilePart(stage)), + Log, + []byte(fmt.Sprintf("failed to get Kueue CR '%s': %v", KueueCRName, err)), + ) + return + } + + payload := map[string]any{ + "stage": stage, + "expected_framework": expectedFramework, + "actual_frameworks": KueueCRFrameworks(kueue), + "available_condition": KueueCRConditionAvailable(kueue), + "cert_condition": KueueCRConditionCertManagerAvailable(kueue), + "name": kueue.GetName(), + "namespace": kueue.GetNamespace(), + } + + data, marshalErr := json.MarshalIndent(payload, "", " ") + if marshalErr != nil { + data = []byte(fmt.Sprintf("failed to marshal diagnostics: %v", marshalErr)) + } + WriteToOutputDir( + test, + fmt.Sprintf("kueue-diagnostics-%s", sanitizeFilePart(stage)), + Log, + data, + ) +} + +func sanitizeFilePart(value string) string { + return strings.ReplaceAll(strings.TrimSpace(value), " ", "-") +} + +func containsString(items []string, target string) bool { + for _, item := range items { + if item == target { + return true + } + } + return false +} diff --git a/tests/kubeflow_sdk/Makefile b/tests/kubeflow_sdk/Makefile new file mode 100644 index 000000000..ee93cd771 --- /dev/null +++ b/tests/kubeflow_sdk/Makefile @@ -0,0 +1,82 @@ +SHELL := /usr/bin/env bash + +SDK_TEST_TIMEOUT ?= 60m +SDK_TEST_EXCLUDE_REGEX ?= +ACCELERATOR_TESTS ?= CUDA +NOTEBOOK_ACCELERATOR ?= +SKIP_S3_TESTS ?= false + +.PHONY: help sdk-test sdk-test-all sdk-test-trainer-all sdk-test-trainer-cpu sdk-test-trainer-sanity sdk-test-trainer-tier1 sdk-test-trainer-cuda sdk-test-trainer-rocm sdk-test-trainer-with-tier sdk-print-selected-trainer-tests sdk-print-trainer-tiers sdk-print-tiers + +help: + @echo "Kubeflow SDK test targets" + @echo "" + @echo "Core:" + @echo " sdk-test Run default SDK set (Sanity + Tier1 + accelerator lane)" + @echo " sdk-test-all Alias for sdk-test (all components, currently trainer only)" + @echo " sdk-test-trainer-all Run all trainer SDK tests (all tiers)" + @echo " sdk-test-trainer-cpu Run trainer SDK CPU lane (Sanity + Tier1)" + @echo " sdk-test-trainer-sanity Run trainer SDK sanity tier" + @echo " sdk-test-trainer-tier1 Run trainer SDK tier1" + @echo " sdk-test-trainer-cuda Run trainer SDK CUDA lane" + @echo " sdk-test-trainer-rocm Run trainer SDK ROCm lane" + @echo " sdk-test-trainer-with-tier Run trainer SDK tests with SDK_TEST_TIER=" + @echo "" + @echo "Debug:" + @echo " sdk-print-selected-trainer-tests" + @echo " sdk-print-trainer-tiers" + @echo " sdk-print-tiers Print tiers for all SDK components" + @echo "" + @echo "Supported env vars:" + @echo " SDK_TEST_TIMEOUT (default: 60m)" + @echo " SDK_TEST_EXCLUDE_REGEX Exclude tests by name regex" + @echo " SKIP_S3_TESTS true|false (exclude S3/AWS-related trainer SDK tests)" + @echo " ACCELERATOR_TESTS CUDA (default) | CPU | ROCM | ALL" + @echo " (controls lane selection and default NOTEBOOK_IMAGE when unset)" + @echo " NOTEBOOK_ACCELERATOR Optional image-only override: CUDA | CPU | ROCM" + @echo " SDK_TEST_TIER Custom trainer tier for sdk-test-trainer-with-tier" + +sdk-test: sdk-test-all + +sdk-test-all: + @set -euo pipefail; \ + $(MAKE) -f tests/kubeflow_sdk/Makefile sdk-test-trainer-sanity SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" ACCELERATOR_TESTS="$(ACCELERATOR_TESTS)" SKIP_S3_TESTS="$(SKIP_S3_TESTS)"; \ + $(MAKE) -f tests/kubeflow_sdk/Makefile sdk-test-trainer-tier1 SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" ACCELERATOR_TESTS="$(ACCELERATOR_TESTS)" SKIP_S3_TESTS="$(SKIP_S3_TESTS)"; \ + case "$$(echo "$(ACCELERATOR_TESTS)" | tr '[:lower:]' '[:upper:]')" in \ + CPU) ;; \ + CUDA) $(MAKE) -f tests/kubeflow_sdk/Makefile sdk-test-trainer-cuda SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" ACCELERATOR_TESTS="$(ACCELERATOR_TESTS)" SKIP_S3_TESTS="$(SKIP_S3_TESTS)" ;; \ + ROCM) $(MAKE) -f tests/kubeflow_sdk/Makefile sdk-test-trainer-rocm SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" ACCELERATOR_TESTS="$(ACCELERATOR_TESTS)" SKIP_S3_TESTS="$(SKIP_S3_TESTS)" ;; \ + ALL) $(MAKE) -f tests/kubeflow_sdk/Makefile sdk-test-trainer-cuda SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" ACCELERATOR_TESTS="$(ACCELERATOR_TESTS)" SKIP_S3_TESTS="$(SKIP_S3_TESTS)"; \ + $(MAKE) -f tests/kubeflow_sdk/Makefile sdk-test-trainer-rocm SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" ACCELERATOR_TESTS="$(ACCELERATOR_TESTS)" SKIP_S3_TESTS="$(SKIP_S3_TESTS)" ;; \ + *) echo "Unsupported ACCELERATOR_TESTS='$(ACCELERATOR_TESTS)'. Use CPU, CUDA, ROCM, or ALL."; exit 1 ;; \ + esac + +sdk-test-trainer-all: + @$(MAKE) -f tests/kubeflow_sdk/trainer/Makefile trainer-sdk-test-all SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" ACCELERATOR_TESTS="$(ACCELERATOR_TESTS)" NOTEBOOK_ACCELERATOR="$(NOTEBOOK_ACCELERATOR)" SKIP_S3_TESTS="$(SKIP_S3_TESTS)" + +sdk-test-trainer-cpu: + @$(MAKE) -f tests/kubeflow_sdk/trainer/Makefile trainer-sdk-test-cpu SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" ACCELERATOR_TESTS="$(ACCELERATOR_TESTS)" NOTEBOOK_ACCELERATOR="CPU" SKIP_S3_TESTS="$(SKIP_S3_TESTS)" + +sdk-test-trainer-sanity: + @$(MAKE) -f tests/kubeflow_sdk/trainer/Makefile trainer-sdk-test-sanity SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" ACCELERATOR_TESTS="$(ACCELERATOR_TESTS)" NOTEBOOK_ACCELERATOR="$(NOTEBOOK_ACCELERATOR)" SKIP_S3_TESTS="$(SKIP_S3_TESTS)" + +sdk-test-trainer-tier1: + @$(MAKE) -f tests/kubeflow_sdk/trainer/Makefile trainer-sdk-test-tier1 SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" ACCELERATOR_TESTS="$(ACCELERATOR_TESTS)" NOTEBOOK_ACCELERATOR="$(NOTEBOOK_ACCELERATOR)" SKIP_S3_TESTS="$(SKIP_S3_TESTS)" + +sdk-test-trainer-cuda: + @$(MAKE) -f tests/kubeflow_sdk/trainer/Makefile trainer-sdk-test-cuda SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" ACCELERATOR_TESTS="$(ACCELERATOR_TESTS)" NOTEBOOK_ACCELERATOR="CUDA" SKIP_S3_TESTS="$(SKIP_S3_TESTS)" + +sdk-test-trainer-rocm: + @$(MAKE) -f tests/kubeflow_sdk/trainer/Makefile trainer-sdk-test-rocm SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" ACCELERATOR_TESTS="$(ACCELERATOR_TESTS)" NOTEBOOK_ACCELERATOR="ROCM" SKIP_S3_TESTS="$(SKIP_S3_TESTS)" + +sdk-test-trainer-with-tier: + @$(MAKE) -f tests/kubeflow_sdk/trainer/Makefile trainer-sdk-test-with-tier SDK_TEST_TIER="$(SDK_TEST_TIER)" SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" ACCELERATOR_TESTS="$(ACCELERATOR_TESTS)" NOTEBOOK_ACCELERATOR="$(NOTEBOOK_ACCELERATOR)" SKIP_S3_TESTS="$(SKIP_S3_TESTS)" + +sdk-print-selected-trainer-tests: + @$(MAKE) -f tests/kubeflow_sdk/trainer/Makefile trainer-sdk-print-selected SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SKIP_S3_TESTS="$(SKIP_S3_TESTS)" + +sdk-print-trainer-tiers: + @$(MAKE) -f tests/kubeflow_sdk/trainer/Makefile trainer-sdk-print-tiers + +sdk-print-tiers: + @$(MAKE) -f tests/kubeflow_sdk/Makefile sdk-print-trainer-tiers diff --git a/tests/kubeflow_sdk/README.md b/tests/kubeflow_sdk/README.md new file mode 100644 index 000000000..7cda00827 --- /dev/null +++ b/tests/kubeflow_sdk/README.md @@ -0,0 +1,425 @@ +# Kubeflow SDK Test Entry Point + +This directory provides a top-level entry point for SDK-focused tests across components. + +For now, it orchestrates existing trainer SDK test entrypoints in `tests/trainer/kubeflow_sdk_test.go` to preserve backward compatibility. + +## Run commands + +From repository root: + +- Show all commands: + - `make -f tests/kubeflow_sdk/Makefile help` +- Run all SDK tests (currently trainer SDK tests): + - `make -f tests/kubeflow_sdk/Makefile sdk-test` +- Run all trainer SDK tests (all tiers): + - `make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-all` +- Run trainer SDK tier shortcuts: + - `make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-cpu` + - `make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-sanity` + - `make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-tier1` + - `make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-cuda` + - `make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-rocm` +- Run trainer SDK with a custom tier: + - `SDK_TEST_TIER=Tier1 make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-with-tier` +- Print selected trainer SDK tests: + - `make -f tests/kubeflow_sdk/Makefile sdk-print-selected-trainer-tests` +- Print trainer tier mapping: + - `make -f tests/kubeflow_sdk/Makefile sdk-print-trainer-tiers` +- Print tiers across all SDK components: + - `make -f tests/kubeflow_sdk/Makefile sdk-print-tiers` + +## Run Examples + +### 1) Trainer sanity tier with a specific SDK git ref + +```bash +NOTEBOOK_USER_TOKEN="$(oc whoami -t)" \ +OPENSHIFT_API_URL="$(oc whoami --show-server)" \ +KUBEFLOW_GIT_URL="kubeflow @ git+https://github.com/opendatahub-io/kubeflow-sdk.git@main" \ +make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-sanity +``` + +What this run does: + +- Executes the trainer sanity tier (`sdk-test-trainer-sanity`). +- Uses the provided token and auto-resolves `NOTEBOOK_USER_NAME` if it is not set. +- Installs Kubeflow SDK from `KUBEFLOW_GIT_URL` (takes precedence over version-based install vars). +- If `NOTEBOOK_IMAGE` is unset, the SDK test harness chooses a default image (CUDA by default unless overridden via `NOTEBOOK_ACCELERATOR`). + +### 2) Tier1 trainer with a specific SDK git ref + +```bash +NOTEBOOK_USER_TOKEN="$(oc whoami -t)" \ +OPENSHIFT_API_URL="$(oc whoami --show-server)" \ +KUBEFLOW_GIT_URL="kubeflow @ git+https://github.com/opendatahub-io/kubeflow-sdk.git@main" \ +make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-tier1 +``` + +What this run does: + +- `sdk-test-trainer-tier1` delegates to the trainer harness target with `SDK_TEST_TIER=Tier1`. +- The harness evaluates `tests/kubeflow_sdk/scripts/install_kubeflow.sh env`, which exports auth/install variables for this run. +- Since `KUBEFLOW_GIT_URL` is set, SDK install is done from that git reference. +- `go test` is run for `./tests/trainer`, and tier tags skip non-`Tier1` tests. + +### 3) Run all default SDK tests (Sanity + Tier1 + accelerator lane (CUDA by default)) + +```bash +NOTEBOOK_USER_TOKEN="$(oc whoami -t)" \ +OPENSHIFT_API_URL="$(oc whoami --show-server)" \ +make -f tests/kubeflow_sdk/Makefile sdk-test +``` + +What this run does by default: + +- Runs `sdk-test-all`, which executes `Sanity`, then `Tier1`, then an accelerator lane. +- The accelerator lane defaults to CUDA because `ACCELERATOR_TESTS` defaults to `CUDA`. +- To change that lane, set `ACCELERATOR_TESTS=CPU`, `ACCELERATOR_TESTS=ROCM`, or `ACCELERATOR_TESTS=ALL`. + +### 4) ROCm lane only (instead of default CUDA) + +```bash +NOTEBOOK_USER_TOKEN="$(oc whoami -t)" \ +OPENSHIFT_API_URL="$(oc whoami --show-server)" \ +ACCELERATOR_TESTS=ROCM \ +make -f tests/kubeflow_sdk/Makefile sdk-test +``` + +What this run does: + +- Runs the default SDK flow with the ROCm accelerator lane instead of CUDA. +- Uses ROCm image defaults unless `NOTEBOOK_IMAGE` is explicitly set. + +### 5) CPU-only lane (Sanity + Tier1, no CUDA/ROCm lane) + +```bash +NOTEBOOK_USER_TOKEN="$(oc whoami -t)" \ +OPENSHIFT_API_URL="$(oc whoami --show-server)" \ +ACCELERATOR_TESTS=CPU \ +make -f tests/kubeflow_sdk/Makefile sdk-test +``` + +What this run does: + +- Runs `Sanity` + `Tier1` and skips CUDA/ROCm accelerator lanes. +- Useful for clusters without GPU capacity. + +### 6) Use username/password instead of token + +```bash +NOTEBOOK_USER_NAME="you@example.com" \ +NOTEBOOK_USER_PASSWORD="***" \ +OPENSHIFT_API_URL="$(oc whoami --show-server)" \ +make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-all +``` + +What this run does: + +- Uses login-based auth to generate a token at runtime. +- Keeps test selection broad by running all trainer SDK tiers. + +### 7) Pin SDK version from custom index + +```bash +NOTEBOOK_USER_TOKEN="$(oc whoami -t)" \ +OPENSHIFT_API_URL="$(oc whoami --show-server)" \ +KUBEFLOW_REQUIRED_VERSION="kubeflow-0.3.0+rhaiv.1" \ +KUBEFLOW_SDK_INDEX_URL="https://console.redhat.com/api/pypi/public-rhai/rhoai/3.4-EA2/cuda13.0-ubi9/simple/" \ +make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-all +``` + +What this run does: + +- Installs a specific SDK version instead of using git source. +- Pulls the package from the specified index URL. + +### 8) Choose notebook image (CUDA, ROCm, CPU, or custom) + +You have four ways to control the notebook image used by SDK runs: + +- **Default behavior (no image vars):** + - Uses accelerator-based defaults from the harness. + - With `sdk-test`, this follows `ACCELERATOR_TESTS` (default: `CUDA`). +- **Select a default image by accelerator only:** + - Set `NOTEBOOK_ACCELERATOR=CUDA|ROCM|CPU`. + - Useful when you want to pick a default image family without hardcoding a full image URL. +- **Drive lane + default image together:** + - Set `ACCELERATOR_TESTS=CUDA|ROCM|CPU|ALL` (for `sdk-test`). + - This changes which lanes run, and also influences default image selection when `NOTEBOOK_IMAGE` is unset. +- **Force a specific image:** + - Set `NOTEBOOK_IMAGE=`. + - This always takes precedence over `NOTEBOOK_ACCELERATOR` and `ACCELERATOR_TESTS`. + +Examples: + +```bash +# Use ROCm default image family for this run +NOTEBOOK_USER_TOKEN="$(oc whoami -t)" \ +OPENSHIFT_API_URL="$(oc whoami --show-server)" \ +NOTEBOOK_ACCELERATOR=ROCM \ +make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-tier1 +``` + +```bash +# Run default suite with CPU lane/default image behavior +NOTEBOOK_USER_TOKEN="$(oc whoami -t)" \ +OPENSHIFT_API_URL="$(oc whoami --show-server)" \ +ACCELERATOR_TESTS=CPU \ +make -f tests/kubeflow_sdk/Makefile sdk-test +``` + +```bash +# Use an explicit custom notebook image (highest priority) +NOTEBOOK_USER_TOKEN="$(oc whoami -t)" \ +OPENSHIFT_API_URL="$(oc whoami --show-server)" \ +NOTEBOOK_IMAGE="quay.io/your-org/your-image@sha256:" \ +make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-all +``` + +### 9) Exclude specific tests by regex + +```bash +NOTEBOOK_USER_TOKEN="$(oc whoami -t)" \ +OPENSHIFT_API_URL="$(oc whoami --show-server)" \ +SDK_TEST_EXCLUDE_REGEX='^(TestA|TestB|TestC|TestD)$' \ +make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-all +``` + +What this run does: + +- Runs the usual selected suite but removes tests whose names match the regex. +- Works with any trainer SDK target (`all`, tiered, or accelerator-specific). + +### 10) Run one specific trainer SDK test by name + +```bash +NOTEBOOK_USER_TOKEN="$(oc whoami -t)" \ +OPENSHIFT_API_URL="$(oc whoami --show-server)" \ +SDK_TRAINER_INCLUDE_REGEX='^TestTorchrunTrainingFailure$' \ +make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-all +``` + +What this run does: + +- Narrows candidate tests to an exact test name match. +- Helpful for fast validation/reproduction while keeping harness setup/auth behavior. + +### 11) Skip all S3/AWS-related trainer SDK tests + +```bash +NOTEBOOK_USER_TOKEN="$(oc whoami -t)" \ +OPENSHIFT_API_URL="$(oc whoami --show-server)" \ +SKIP_S3_TESTS=true \ +make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-all +``` + +What this run does: + +- Automatically excludes S3/AWS-related trainer tests (currently those matching `^TestRhaiS3`). +- Combines with `SDK_TEST_EXCLUDE_REGEX` if you set both. + +## Writing logs with `TEST_OUTPUT_DIR` + +Use `TEST_OUTPUT_DIR` when you want test artifacts saved under a stable local directory instead of ephemeral temp paths. + +How it works: + +- Set `TEST_OUTPUT_DIR` to a parent directory (absolute path recommended). +- The test framework creates a per-test subdirectory under that parent (it does not write all runs into one flat folder). +- During test cleanup, pod container logs and namespace event logs are written into that output directory. +- If `TEST_OUTPUT_DIR` is unset, an ephemeral temp directory is used and its path is only visible in test output logs. + +Example: + +```bash +mkdir -p /tmp/sdk-test-logs +NOTEBOOK_USER_TOKEN="$(oc whoami -t)" \ +OPENSHIFT_API_URL="$(oc whoami --show-server)" \ +TEST_OUTPUT_DIR=/tmp/sdk-test-logs \ +make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-all +``` + +## Tier-aware runs + +You can run predefined trainer tier targets: + +- `make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-cpu` +- `make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-sanity` +- `make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-tier1` +- `make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-cuda` +- `make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-rocm` + +Trainer SDK tiers exposed by this harness: + +- `Sanity` +- `Tier1` +- `CPU` (harness lane composed of `Sanity` + `Tier1`) +- `CUDA` (internally maps to trainer `KFTO-CUDA`) +- `ROCm` (internally maps to trainer `KFTO-ROCm`) + +Default `sdk-test` behavior: + +- runs `Sanity` +- runs `Tier1` +- runs accelerator lane from `ACCELERATOR_TESTS` (default `CUDA`; supported: `CPU`, `CUDA`, `ROCM`, `ALL`) +- if `NOTEBOOK_IMAGE` is unset, the SDK test harness picks a default image by accelerator: + - `CUDA`/`ALL` -> `quay.io/opendatahub/odh-training-cuda128-torch29-py312@sha256:87539ef75e399efceefc6ecc54ebdc4453f794302f417bd30372112692eee70c` + - `ROCM` -> `quay.io/opendatahub/odh-training-rocm64-torch29-py312:odh-stable` + - `CPU` -> `quay.io/rhoai/odh-workbench-jupyter-datascience-cpu-py312-rhel9:rhoai-3.4` +- you can override image selection only (without changing test lane) via `NOTEBOOK_ACCELERATOR=CUDA|ROCM|CPU` + +## Excluding tests + +Use `SDK_TEST_EXCLUDE_REGEX` to exclude by test name: + +- Exclude a single test: + - `SDK_TEST_EXCLUDE_REGEX='TestTorchrunTrainingFailure' make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-all` +- Exclude a group: + - `SDK_TEST_EXCLUDE_REGEX='TestRhaiS3' make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-all` +- Exclude while using a tier: + - `SDK_TEST_EXCLUDE_REGEX='TestRhaiS3' make -f tests/kubeflow_sdk/Makefile sdk-test-trainer-tier1` + +## Environment variable reference + +The SDK test harness accepts the following environment variables. + +### Authentication and cluster access + +- `NOTEBOOK_USER_TOKEN` + - Preferred auth input for SDK runs. + - If set, the harness uses it directly and derives `NOTEBOOK_USER_NAME` via `oc whoami --token=...` when needed. +- `NOTEBOOK_USER_NAME` + - Username for login-based auth mode. + - Required with `NOTEBOOK_USER_PASSWORD` when `NOTEBOOK_USER_TOKEN` is not provided. +- `NOTEBOOK_USER_PASSWORD` + - Password paired with `NOTEBOOK_USER_NAME` for `oc login` token generation. +- `OPENSHIFT_API_URL` + - OpenShift API endpoint used for `oc login` / `oc whoami --token`. + - If unset, the harness attempts to resolve it with `oc whoami --show-server`. + +Auth precedence: + +1. `NOTEBOOK_USER_TOKEN` +2. `NOTEBOOK_USER_NAME` + `NOTEBOOK_USER_PASSWORD` + +### SDK installation source + +- `KUBEFLOW_GIT_URL` + - Git install source, for example: + - `kubeflow @ git+https://github.com//.git@` + - If set, the harness forces git-based install (`KUBEFLOW_INSTALL_FROM_GIT=true` internally). +- `KUBEFLOW_REQUIRED_VERSION` + - Version to install when git URL is not provided. + - Example: `0.2.1+rhai2` +- `KUBEFLOW_SDK_INDEX_URL` + - Optional package index URL for version-based installs. + - Exported to installer as `KUBEFLOW_PYPI_INDEX_URL`. + +Install precedence: + +1. `KUBEFLOW_GIT_URL` +2. `KUBEFLOW_REQUIRED_VERSION` (+ optional `KUBEFLOW_SDK_INDEX_URL`) +3. Installer default behavior in `install_kubeflow.py` + +### Notebook image and accelerator selection + +- `NOTEBOOK_IMAGE` + - Explicit notebook image override. + - If set, this always wins over image defaults/accelerator-based selection. +- `ACCELERATOR_TESTS` + - Controls lane selection for `sdk-test`. + - Supported values: `CUDA` (default), `CPU`, `ROCM`, `ALL`. + - Also used as default image selector when `NOTEBOOK_IMAGE` is unset. +- `NOTEBOOK_ACCELERATOR` + - Optional image-only selector (`CUDA|ROCM|CPU`) when `NOTEBOOK_IMAGE` is unset. + - Does not change tier/lane logic by itself. + +Default image mapping (used only when `NOTEBOOK_IMAGE` is unset): + +- `CUDA`/`ALL` -> `quay.io/opendatahub/odh-training-cuda128-torch29-py312@sha256:87539ef75e399efceefc6ecc54ebdc4453f794302f417bd30372112692eee70c` +- `ROCM` -> `quay.io/opendatahub/odh-training-rocm64-torch29-py312:odh-stable` +- `CPU` -> `quay.io/rhoai/odh-workbench-jupyter-datascience-cpu-py312-rhel9:rhoai-3.4` + +### Test selection and execution controls + +- `SDK_TEST_TIMEOUT` + - `go test` timeout passed to SDK harness targets. + - Default: `60m`. +- `SDK_TEST_EXCLUDE_REGEX` + - Regex for excluding test names from selected SDK trainer tests. + - Applied after include selection and tier filtering. +- `SKIP_S3_TESTS` + - Convenience switch to exclude all S3/AWS-related trainer SDK tests. + - Supported values: `true` or `false` (default: `false`). + - Under the hood this appends an S3 exclusion regex to the active exclude filter. +- `SDK_TEST_TIER` + - Tier value used by `sdk-test-trainer-with-tier`. + - Typical values include `Sanity`, `Tier1`, `KFTO-CUDA`, `KFTO-ROCm`. +- `SDK_TRAINER_INCLUDE_REGEX` (advanced) + - Overrides trainer SDK include pattern used to enumerate candidate tests. + - Default pattern maps to test names in `tests/trainer/kubeflow_sdk_test.go`. +- `SDK_TRAINER_SANITY_INCLUDE_REGEX` (advanced) + - Overrides the internal sanity include pattern used by trainer harness internals. + - Default: `^TestKubeflowSdk`. +- `SDK_TEST_S3_EXCLUDE_REGEX` (advanced) + - Trainer-level regex used when `SKIP_S3_TESTS=true`. + - Default: `^TestRhaiS3`. +- `TEST_OUTPUT_DIR` + - Optional output directory for test artifacts/logs when supported by underlying test helpers. + - If unset, test framework may use a temporary directory. + +### Storage and dataset inputs (trainer SDK tests) + +These are needed for trainer SDK scenarios that read/write artifacts from object storage. + +Where these are used: + +- Trainer SDK test entrypoints: `tests/trainer/kubeflow_sdk_test.go` +- S3-focused trainer tests include: + - `TestRhaiS3CheckpointingCPU` + - `TestRhaiS3FsdpFullStateCheckpointingCPU` + - `TestRhaiS3FsdpSharedStateCheckpointingCuda` + - `TestRhaiS3DeepspeedStage0CheckpointingCuda` + +- `AWS_DEFAULT_ENDPOINT` + - S3-compatible endpoint URL/host used by notebook jobs. +- `AWS_ACCESS_KEY_ID` + - Access key for object storage authentication. +- `AWS_SECRET_ACCESS_KEY` + - Secret key for object storage authentication. +- `AWS_STORAGE_BUCKET` + - Bucket name used for datasets/checkpoints/artifacts. +- `AWS_DEFAULT_REGION` + - Optional region used by some provider/client paths. +- `AWS_STORAGE_BUCKET_MNIST_DIR` + - Prefix/path for MNIST-related data. +- `AWS_STORAGE_BUCKET_OSFT_DIR` + - Prefix/path for OSFT-related data. +- `AWS_STORAGE_BUCKET_SFT_DIR` + - Prefix/path for SFT-related data. +- `AWS_STORAGE_BUCKET_LORA_DIR` + - Prefix/path for LoRA-related data. +- `MODEL_S3_PREFIX` + - Optional model prefix used by selected RHAI scenarios. +- `DATASET_S3_PREFIX` + - Optional dataset prefix used by selected RHAI scenarios. + +Other optional trainer test controls: + +- `TEST_NAMESPACE_NAME` + - Reuse an existing namespace instead of creating a generated test namespace. + +### Notes + +- `KUBEFLOW_INSTALL_FROM_GIT` is exported internally by the harness when `KUBEFLOW_GIT_URL` is provided; you do not need to set it for harness-based runs. +- For direct (non-harness) trainer invocations, behavior may differ because those commands bypass `tests/kubeflow_sdk/scripts/install_kubeflow.sh`. + +## Scaling to future components + +Future SDK components can add folders like: + +- `tests/kubeflow_sdk/core` +- `tests/kubeflow_sdk/spark` + +and extend the Makefile with component-specific targets while reusing shared auth/install env preparation in `tests/kubeflow_sdk/scripts/install_kubeflow.sh`. diff --git a/tests/kubeflow_sdk/scripts/install_kubeflow.sh b/tests/kubeflow_sdk/scripts/install_kubeflow.sh new file mode 100755 index 000000000..c3934bf03 --- /dev/null +++ b/tests/kubeflow_sdk/scripts/install_kubeflow.sh @@ -0,0 +1,119 @@ +#!/usr/bin/env bash +set -euo pipefail + +mode="${1:-env}" +DEFAULT_NOTEBOOK_IMAGE_CUDA="quay.io/opendatahub/odh-training-cuda128-torch29-py312@sha256:87539ef75e399efceefc6ecc54ebdc4453f794302f417bd30372112692eee70c" +DEFAULT_NOTEBOOK_IMAGE_ROCM="quay.io/opendatahub/odh-training-rocm64-torch29-py312:odh-stable" +DEFAULT_NOTEBOOK_IMAGE_CPU="quay.io/rhoai/odh-workbench-jupyter-datascience-cpu-py312-rhel9:rhoai-3.4" + +if [[ "${mode}" != "env" ]]; then + echo "Usage: $0 env" >&2 + exit 2 +fi + +trim() { + local value="${1:-}" + # shellcheck disable=SC2001 + echo "$(echo "${value}" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')" +} + +resolve_api_server() { + if [[ -n "${OPENSHIFT_API_URL:-}" ]]; then + echo "${OPENSHIFT_API_URL}" + return + fi + oc whoami --show-server +} + +resolve_token_from_userpass() { + local user="${1}" + local password="${2}" + local api_server + api_server="$(resolve_api_server)" + local tmp_kubeconfig + tmp_kubeconfig="$(mktemp)" + trap 'rm -f "${tmp_kubeconfig}"' RETURN + + oc login -u "${user}" -p "${password}" "${api_server}" --insecure-skip-tls-verify=true --kubeconfig="${tmp_kubeconfig}" >/dev/null + oc whoami --show-token --kubeconfig="${tmp_kubeconfig}" +} + +resolve_username_from_token() { + local token="${1}" + local api_server + api_server="$(resolve_api_server)" + oc whoami --token="${token}" --server="${api_server}" --insecure-skip-tls-verify=true +} + +resolve_default_notebook_image() { + local accelerator_input="${1:-CUDA}" + local accelerator + accelerator="$(echo "${accelerator_input}" | tr '[:lower:]' '[:upper:]')" + case "${accelerator}" in + ROCM) + echo "${DEFAULT_NOTEBOOK_IMAGE_ROCM}" + ;; + CPU) + echo "${DEFAULT_NOTEBOOK_IMAGE_CPU}" + ;; + CUDA|ALL|"") + echo "${DEFAULT_NOTEBOOK_IMAGE_CUDA}" + ;; + *) + # Unknown selector: keep behavior predictable and default to CUDA. + echo "${DEFAULT_NOTEBOOK_IMAGE_CUDA}" + ;; + esac +} + +token="$(trim "${NOTEBOOK_USER_TOKEN:-}")" +username="$(trim "${NOTEBOOK_USER_NAME:-}")" +password="$(trim "${NOTEBOOK_USER_PASSWORD:-}")" + +if [[ -z "${token}" ]]; then + if [[ -z "${username}" || -z "${password}" ]]; then + echo "Missing auth input. Provide NOTEBOOK_USER_TOKEN or NOTEBOOK_USER_NAME + NOTEBOOK_USER_PASSWORD." >&2 + exit 1 + fi + token="$(resolve_token_from_userpass "${username}" "${password}")" +fi + +if [[ -z "${username}" ]]; then + username="$(resolve_username_from_token "${token}")" +fi + +if [[ -z "${username}" ]]; then + echo "Unable to resolve NOTEBOOK_USER_NAME." >&2 + exit 1 +fi + +git_url="$(trim "${KUBEFLOW_GIT_URL:-}")" +version="$(trim "${KUBEFLOW_REQUIRED_VERSION:-}")" +index_url="$(trim "${KUBEFLOW_SDK_INDEX_URL:-}")" +notebook_image="$(trim "${NOTEBOOK_IMAGE:-}")" +accelerator_tests="$(trim "${ACCELERATOR_TESTS:-}")" +notebook_accelerator="$(trim "${NOTEBOOK_ACCELERATOR:-}")" +if [[ -z "${notebook_accelerator}" ]]; then + notebook_accelerator="${accelerator_tests}" +fi + +echo "export NOTEBOOK_USER_TOKEN='${token}'" +echo "export NOTEBOOK_USER_NAME='${username}'" + +if [[ -n "${notebook_image}" ]]; then + echo "export NOTEBOOK_IMAGE='${notebook_image}'" +else + default_notebook_image="$(resolve_default_notebook_image "${notebook_accelerator}")" + echo "export NOTEBOOK_IMAGE='${default_notebook_image}'" +fi + +if [[ -n "${git_url}" ]]; then + echo "export KUBEFLOW_INSTALL_FROM_GIT='true'" + echo "export KUBEFLOW_GIT_URL='${git_url}'" +elif [[ -n "${version}" ]]; then + echo "export KUBEFLOW_REQUIRED_VERSION='${version}'" +fi + +if [[ -n "${index_url}" ]]; then + echo "export KUBEFLOW_PYPI_INDEX_URL='${index_url}'" +fi diff --git a/tests/kubeflow_sdk/trainer/Makefile b/tests/kubeflow_sdk/trainer/Makefile new file mode 100644 index 000000000..37aa31aea --- /dev/null +++ b/tests/kubeflow_sdk/trainer/Makefile @@ -0,0 +1,110 @@ +SHELL := /usr/bin/env bash + +SDK_TEST_TIMEOUT ?= 60m +SDK_TEST_EXCLUDE_REGEX ?= +SDK_TEST_TIER ?= +ACCELERATOR_TESTS ?= CUDA +NOTEBOOK_ACCELERATOR ?= +SKIP_S3_TESTS ?= false +SDK_TEST_S3_EXCLUDE_REGEX ?= ^TestRhaiS3 + +RG_BIN := $(shell command -v rg 2>/dev/null) +ifeq ($(RG_BIN),) + TEST_MATCH_CMD := grep -E + TEST_EXCLUDE_CMD := grep -Ev +else + TEST_MATCH_CMD := rg + TEST_EXCLUDE_CMD := rg -v +endif + +# Source of truth remains tests/trainer/kubeflow_sdk_test.go +SDK_TRAINER_INCLUDE_REGEX ?= ^(TestKubeflowSdk|TestOsftTrainingHub|TestLoraTrainingHub|TestSftTrainingHub|TestRhai|TestTrainingFailureScenarios|TestTorchrunTrainingFailure) +SDK_TRAINER_SANITY_INCLUDE_REGEX ?= ^TestKubeflowSdk + +.PHONY: trainer-sdk-test trainer-sdk-test-with-tier trainer-sdk-test-all trainer-sdk-test-cpu trainer-sdk-test-sanity trainer-sdk-test-tier1 trainer-sdk-test-cuda trainer-sdk-test-rocm trainer-sdk-test-name-sanity trainer-sdk-print-selected trainer-sdk-print-tiers + +trainer-sdk-test: + @set -euo pipefail; \ + eval "$$(ACCELERATOR_TESTS='$(ACCELERATOR_TESTS)' NOTEBOOK_ACCELERATOR='$(NOTEBOOK_ACCELERATOR)' ./tests/kubeflow_sdk/scripts/install_kubeflow.sh env)"; \ + EFFECTIVE_EXCLUDE_REGEX="$${SDK_TEST_EXCLUDE_REGEX:-}"; \ + if [[ "$$(echo "$(SKIP_S3_TESTS)" | tr '[:lower:]' '[:upper:]')" == "TRUE" ]]; then \ + if [[ -n "$$EFFECTIVE_EXCLUDE_REGEX" ]]; then \ + EFFECTIVE_EXCLUDE_REGEX="($$EFFECTIVE_EXCLUDE_REGEX)|($(SDK_TEST_S3_EXCLUDE_REGEX))"; \ + else \ + EFFECTIVE_EXCLUDE_REGEX="$(SDK_TEST_S3_EXCLUDE_REGEX)"; \ + fi; \ + fi; \ + TESTS="$$(TEST_TIER="$(SDK_TEST_TIER)" go test ./tests/trainer -list '$(SDK_TRAINER_INCLUDE_REGEX)' | $(TEST_MATCH_CMD) '^Test' | (if [[ -n "$$EFFECTIVE_EXCLUDE_REGEX" ]]; then $(TEST_EXCLUDE_CMD) "$$EFFECTIVE_EXCLUDE_REGEX"; else awk '1'; fi) | paste -sd'|' -)"; \ + if [[ -z "$$TESTS" ]]; then \ + echo "No SDK trainer tests matched include/exclude filters."; \ + exit 1; \ + fi; \ + if [[ "$$(echo "$(SKIP_S3_TESTS)" | tr '[:lower:]' '[:upper:]')" == "TRUE" ]]; then \ + echo "SKIP_S3_TESTS=true active. Excluding tests matching: $(SDK_TEST_S3_EXCLUDE_REGEX)"; \ + fi; \ + echo "Running SDK trainer tests (TEST_TIER='$(SDK_TEST_TIER)'): $$TESTS"; \ + TEST_TIER="$(SDK_TEST_TIER)" go test -timeout "$(SDK_TEST_TIMEOUT)" -count=1 ./tests/trainer -run "^($$TESTS)$$" -v + +trainer-sdk-test-with-tier: + @if [[ -z "$(SDK_TEST_TIER)" ]]; then \ + echo "SDK_TEST_TIER is required (example: SDK_TEST_TIER=Tier1)"; \ + exit 1; \ + fi + @$(MAKE) -f tests/kubeflow_sdk/trainer/Makefile trainer-sdk-test SDK_TEST_TIER="$(SDK_TEST_TIER)" SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" + +trainer-sdk-test-all: + @$(MAKE) -f tests/kubeflow_sdk/trainer/Makefile trainer-sdk-test SDK_TEST_TIER="" SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" + +trainer-sdk-test-cpu: + @$(MAKE) -f tests/kubeflow_sdk/trainer/Makefile trainer-sdk-test-sanity SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" + @$(MAKE) -f tests/kubeflow_sdk/trainer/Makefile trainer-sdk-test-tier1 SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" + +trainer-sdk-test-sanity: + @$(MAKE) -f tests/kubeflow_sdk/trainer/Makefile trainer-sdk-test-name-sanity SDK_TEST_TIER="Sanity" SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" + +trainer-sdk-test-tier1: + @$(MAKE) -f tests/kubeflow_sdk/trainer/Makefile trainer-sdk-test SDK_TEST_TIER="Tier1" SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" + +trainer-sdk-test-cuda: + @$(MAKE) -f tests/kubeflow_sdk/trainer/Makefile trainer-sdk-test SDK_TEST_TIER="KFTO-CUDA" SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" + +trainer-sdk-test-rocm: + @$(MAKE) -f tests/kubeflow_sdk/trainer/Makefile trainer-sdk-test SDK_TEST_TIER="KFTO-ROCm" SDK_TEST_EXCLUDE_REGEX="$(SDK_TEST_EXCLUDE_REGEX)" SDK_TEST_TIMEOUT="$(SDK_TEST_TIMEOUT)" + +trainer-sdk-test-name-sanity: + @set -euo pipefail; \ + eval "$$(ACCELERATOR_TESTS='$(ACCELERATOR_TESTS)' NOTEBOOK_ACCELERATOR='$(NOTEBOOK_ACCELERATOR)' ./tests/kubeflow_sdk/scripts/install_kubeflow.sh env)"; \ + EFFECTIVE_EXCLUDE_REGEX="$${SDK_TEST_EXCLUDE_REGEX:-}"; \ + if [[ "$$(echo "$(SKIP_S3_TESTS)" | tr '[:lower:]' '[:upper:]')" == "TRUE" ]]; then \ + if [[ -n "$$EFFECTIVE_EXCLUDE_REGEX" ]]; then \ + EFFECTIVE_EXCLUDE_REGEX="($$EFFECTIVE_EXCLUDE_REGEX)|($(SDK_TEST_S3_EXCLUDE_REGEX))"; \ + else \ + EFFECTIVE_EXCLUDE_REGEX="$(SDK_TEST_S3_EXCLUDE_REGEX)"; \ + fi; \ + fi; \ + TESTS="$$(TEST_TIER="$(SDK_TEST_TIER)" go test ./tests/trainer -list '$(SDK_TRAINER_SANITY_INCLUDE_REGEX)' | $(TEST_MATCH_CMD) '^Test' | (if [[ -n "$$EFFECTIVE_EXCLUDE_REGEX" ]]; then $(TEST_EXCLUDE_CMD) "$$EFFECTIVE_EXCLUDE_REGEX"; else awk '1'; fi) | paste -sd'|' -)"; \ + if [[ -z "$$TESTS" ]]; then \ + echo "No SDK sanity tests matched include/exclude filters."; \ + exit 1; \ + fi; \ + echo "Running SDK sanity tests (TEST_TIER='$(SDK_TEST_TIER)'): $$TESTS"; \ + TEST_TIER="$(SDK_TEST_TIER)" go test -timeout "$(SDK_TEST_TIMEOUT)" -count=1 ./tests/trainer -run "^($$TESTS)$$" -v + +trainer-sdk-print-selected: + @set -euo pipefail; \ + EFFECTIVE_EXCLUDE_REGEX="$${SDK_TEST_EXCLUDE_REGEX:-}"; \ + if [[ "$$(echo "$(SKIP_S3_TESTS)" | tr '[:lower:]' '[:upper:]')" == "TRUE" ]]; then \ + if [[ -n "$$EFFECTIVE_EXCLUDE_REGEX" ]]; then \ + EFFECTIVE_EXCLUDE_REGEX="($$EFFECTIVE_EXCLUDE_REGEX)|($(SDK_TEST_S3_EXCLUDE_REGEX))"; \ + else \ + EFFECTIVE_EXCLUDE_REGEX="$(SDK_TEST_S3_EXCLUDE_REGEX)"; \ + fi; \ + fi; \ + TEST_TIER="$(SDK_TEST_TIER)" go test ./tests/trainer -list '$(SDK_TRAINER_INCLUDE_REGEX)' | $(TEST_MATCH_CMD) '^Test' | (if [[ -n "$$EFFECTIVE_EXCLUDE_REGEX" ]]; then $(TEST_EXCLUDE_CMD) "$$EFFECTIVE_EXCLUDE_REGEX"; else awk '1'; fi) + +trainer-sdk-print-tiers: + @echo "Trainer SDK tiers exposed via this wrapper:" + @echo " Sanity" + @echo " Tier1" + @echo " CUDA (mapped to trainer KFTO-CUDA)" + @echo " ROCm (mapped to trainer KFTO-ROCm)" diff --git a/tests/kubeflow_sdk/trainer/README.md b/tests/kubeflow_sdk/trainer/README.md new file mode 100644 index 000000000..88c665951 --- /dev/null +++ b/tests/kubeflow_sdk/trainer/README.md @@ -0,0 +1,25 @@ +# Trainer SDK Wrapper + +This folder documents the trainer component entry from the top-level SDK test harness. + +Current model: + +- SDK test entrypoints remain in `tests/trainer/kubeflow_sdk_test.go`. +- `tests/kubeflow_sdk/Makefile` exposes top-level SDK commands. +- `tests/kubeflow_sdk/trainer/Makefile` contains trainer-specific execution logic. + +This keeps existing `go test ./tests/trainer ...` workflows intact while adding a dedicated SDK-centric entry point. + +## Current trainer SDK test selection + +The wrapper includes tests matching: + +- `TestKubeflowSdk*` +- `TestOsftTrainingHub*` +- `TestLoraTrainingHub*` +- `TestSftTrainingHub*` +- `TestRhai*` +- `TestTrainingFailureScenarios` +- `TestTorchrunTrainingFailure` + +Use `SDK_TEST_EXCLUDE_REGEX` to omit specific tests by name pattern. diff --git a/tests/trainer/resources/disconnected_env/install_kubeflow.py b/tests/trainer/resources/disconnected_env/install_kubeflow.py index 826927b21..7a28fc919 100644 --- a/tests/trainer/resources/disconnected_env/install_kubeflow.py +++ b/tests/trainer/resources/disconnected_env/install_kubeflow.py @@ -48,6 +48,10 @@ def get_rhai_pypi_index() -> str: - CUDA: https://console.redhat.com/api/pypi/public-rhai/rhoai/3.3/cuda12.9-ubi9/simple/ - ROCm: https://console.redhat.com/api/pypi/public-rhai/rhoai/3.3/rocm6.4-ubi9/simple/ """ + custom_index = os.environ.get("KUBEFLOW_SDK_INDEX_URL") + if custom_index: + return custom_index + gpu_type = os.environ.get("GPU_TYPE", "cpu").lower() base = "https://console.redhat.com/api/pypi/public-rhai/rhoai/3.3" diff --git a/tests/trainer/resources/failure_scenarios.ipynb b/tests/trainer/resources/failure_scenarios.ipynb index fce80dc78..b38dba47d 100644 --- a/tests/trainer/resources/failure_scenarios.ipynb +++ b/tests/trainer/resources/failure_scenarios.ipynb @@ -1,132 +1,410 @@ { - "cells": [ - { - "cell_type": "code", - "id": "cell-pip-install", - "metadata": {}, - "outputs": [], - "source": [ - "# pip Install kubeflow SDK from main branch for testing\n", - "%pip install git+https://github.com/opendatahub-io/kubeflow-sdk.git@main" - ], - "execution_count": null + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "cell-pip-install", + "metadata": {}, + "outputs": [], + "source": [ + "# kubeflow SDK is installed by test harness via install_kubeflow.py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-k8s-setup", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import urllib3\n", + "from kubernetes import client as k8s\n", + "\n", + "# Suppress InsecureRequestWarning since we use verify_ssl = False\n", + "urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)\n", + "\n", + "api_server = os.getenv(\"OPENSHIFT_API_URL\")\n", + "token = os.getenv(\"NOTEBOOK_USER_TOKEN\")\n", + "if not api_server or not token:\n", + " raise RuntimeError(\"OPENSHIFT_API_URL and NOTEBOOK_USER_TOKEN environment variables are required\")\n", + "\n", + "configuration = k8s.Configuration()\n", + "configuration.host = api_server\n", + "configuration.verify_ssl = False\n", + "configuration.api_key = {\"authorization\": f\"Bearer {token}\"}\n", + "api_client = k8s.ApiClient(configuration)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-trainer-client", + "metadata": {}, + "outputs": [], + "source": [ + "from kubeflow.trainer import TrainerClient\n", + "from kubeflow.common.types import KubernetesBackendConfig\n", + "\n", + "backend_cfg = KubernetesBackendConfig(\n", + " client_configuration=api_client.configuration,\n", + ")\n", + "\n", + "client = TrainerClient(backend_cfg)\n", + "print(client)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-find-runtime", + "metadata": {}, + "outputs": [], + "source": [ + "training_runtime_name = os.getenv(\"TRAINING_RUNTIME\")\n", + "if not training_runtime_name:\n", + " raise RuntimeError(\"TRAINING_RUNTIME environment variable is required\")\n", + "\n", + "th_runtime = None\n", + "for runtime in client.list_runtimes():\n", + " if runtime.name == training_runtime_name:\n", + " th_runtime = runtime\n", + " print(\"Found runtime: \" + str(th_runtime))\n", + " break\n", + "\n", + "if th_runtime is None:\n", + " raise RuntimeError(f\"Required runtime '{training_runtime_name}' not found\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-helper", + "metadata": {}, + "outputs": [], + "source": [ + "import time\n", + "\n", + "from kubeflow.trainer.rhai import TrainingHubAlgorithms, TrainingHubTrainer\n", + "\n", + "NAMESPACE = os.getenv(\"NOTEBOOK_NAMESPACE\")\n", + "print(f\"Using namespace: {NAMESPACE}\")\n", + "\n", + "# Event reasons that indicate a container crash or failure\n", + "FAILURE_EVENT_REASONS = {\"BackOff\", \"CrashLoopBackOff\", \"Failed\", \"OOMKilled\", \"OOMKilling\", \"Killing\"}\n", + "\n", + "\n", + "def check_job_failure(client, job_name):\n", + " \"\"\"Check if a TrainJob has failed using SDK APIs (get_job and get_job_events).\n", + "\n", + " Returns (failure_confirmed, details) where failure is confirmed if:\n", + " - get_job().status == \"Failed\", OR\n", + " - get_job_events() contains crash-related events (BackOff, OOMKilled, etc.)\n", + " \"\"\"\n", + " failure_confirmed = False\n", + " details = []\n", + "\n", + " # Check job status via get_job()\n", + " try:\n", + " job = client.get_job(name=job_name)\n", + " if job.status == \"Failed\":\n", + " failure_confirmed = True\n", + " details.append(f\"job.status=Failed\")\n", + " except Exception as e:\n", + " print(f\" get_job() error: {e}\")\n", + "\n", + " # Check events via get_job_events()\n", + " try:\n", + " events = client.get_job_events(name=job_name)\n", + " for event in events:\n", + " reason = getattr(event, \"reason\", \"\") or \"\"\n", + " if reason in FAILURE_EVENT_REASONS:\n", + " failure_confirmed = True\n", + " message = getattr(event, \"message\", \"\") or \"\"\n", + " details.append(f\"event: reason={reason} message={message[:100]}\")\n", + " except Exception as e:\n", + " print(f\" get_job_events() error: {e}\")\n", + "\n", + " return failure_confirmed, details\n", + "\n", + "\n", + "def run_failure_scenario(client, scenario_name, training_params, expected_error, runtime, algorithm=TrainingHubAlgorithms.SFT):\n", + " \"\"\"Submit a TrainJob expected to fail, verify the expected error in logs\n", + " and that the job failure is detected via get_job()/get_job_events().\"\"\"\n", + " print(f\"\\n{'='*60}\")\n", + " print(f\"Scenario: {scenario_name}\")\n", + " print(f\"Expected error: {expected_error}\")\n", + " print(f\"{'='*60}\")\n", + "\n", + " job_name = None\n", + "\n", + " try:\n", + " job_name = client.train(\n", + " trainer=TrainingHubTrainer(\n", + " algorithm=algorithm,\n", + " func_args=training_params,\n", + " resources_per_node={\n", + " \"cpu\": 4,\n", + " \"memory\": \"16Gi\",\n", + " },\n", + " ),\n", + " runtime=runtime,\n", + " )\n", + " print(f\"TrainJob created: {job_name}\")\n", + "\n", + " # Wait for the job to start running\n", + " try:\n", + " client.wait_for_job_status(name=job_name, status={\"Running\", \"Failed\"}, timeout=300)\n", + " except Exception as e:\n", + " print(f\"Wait for Running/Failed raised: {e}\")\n", + "\n", + " # Poll until we see: (1) the expected error in logs, and (2) failure via SDK APIs\n", + " error_found = False\n", + " failure_confirmed = False\n", + " deadline = time.time() + 300 # 5 minute timeout\n", + "\n", + " while time.time() < deadline:\n", + " # Check logs for the expected error\n", + " if not error_found:\n", + " try:\n", + " log_lines = list(client.get_job_logs(job_name, follow=False))\n", + " log_text = \"\\n\".join(str(line) for line in log_lines)\n", + " if expected_error in log_text:\n", + " error_found = True\n", + " print(f\"Found expected error in logs: '{expected_error}'\")\n", + " except Exception as e:\n", + " print(f\"Log fetch error (retrying): {e}\")\n", + "\n", + " # Check job failure via get_job() and get_job_events()\n", + " if not failure_confirmed:\n", + " failure_confirmed, details = check_job_failure(client, job_name)\n", + " if failure_confirmed:\n", + " print(f\"Job failure confirmed via SDK: {'; '.join(details)}\")\n", + "\n", + " if error_found and failure_confirmed:\n", + " break\n", + "\n", + " time.sleep(15)\n", + "\n", + " # Verify results\n", + " assert error_found, f\"Expected error '{expected_error}' not found in logs\"\n", + " assert failure_confirmed, (\n", + " f\"Job failure not confirmed via get_job()/get_job_events() \"\n", + " f\"— final job status: {client.get_job(name=job_name).status}\"\n", + " )\n", + "\n", + " print(f\"PASSED: {scenario_name}\")\n", + " result = True\n", + "\n", + " except Exception as e:\n", + " print(f\"FAILED: {scenario_name} - {e}\")\n", + " result = False\n", + "\n", + " finally:\n", + " # Cleanup\n", + " if job_name:\n", + " try:\n", + " client.delete_job(job_name)\n", + " print(f\"Deleted job: {job_name}\")\n", + " except Exception as e:\n", + " print(f\"Warning: failed to delete job {job_name}: {e}\")\n", + "\n", + " return result" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-scenario-invalid-dataset", + "metadata": {}, + "outputs": [], + "source": [ + "# Scenario 1: Invalid Dataset Path\n", + "result_1 = run_failure_scenario(\n", + " client=client,\n", + " scenario_name=\"Invalid Dataset Path\",\n", + " training_params={\n", + " \"model_path\": \"/opt/app-root/src\",\n", + " \"data_path\": \"/nonexistent/dataset.jsonl\",\n", + " \"ckpt_output_dir\": \"/opt/app-root/src/checkpoints\",\n", + " \"data_output_dir\": \"/opt/app-root/src/sft-data\",\n", + " \"effective_batch_size\": 128,\n", + " \"num_epochs\": 1,\n", + " \"max_seq_len\": 4096,\n", + " \"max_batch_len\": 10000,\n", + " \"learning_rate\": 5e-6,\n", + " },\n", + " expected_error=\"FileNotFoundError\",\n", + " runtime=th_runtime,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-scenario-invalid-model", + "metadata": {}, + "outputs": [], + "source": [ + "# Scenario 2: Invalid Model Path\n", + "result_2 = run_failure_scenario(\n", + " client=client,\n", + " scenario_name=\"Invalid Model Path\",\n", + " training_params={\n", + " \"model_path\": \"/nonexistent/model\",\n", + " \"data_path\": \"/opt/app-root/src/data.jsonl\",\n", + " \"ckpt_output_dir\": \"/opt/app-root/src/checkpoints\",\n", + " \"data_output_dir\": \"/opt/app-root/src/sft-data\",\n", + " \"effective_batch_size\": 128,\n", + " \"num_epochs\": 1,\n", + " \"max_seq_len\": 4096,\n", + " \"max_batch_len\": 10000,\n", + " \"learning_rate\": 5e-6,\n", + " },\n", + " expected_error=\"FileNotFoundError\",\n", + " runtime=th_runtime,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-scenario-invalid-hyperparams", + "metadata": {}, + "outputs": [], + "source": [ + "# Scenario 3: Invalid Training Hyperparameters\n", + "# Uses a string value for max_batch_len where an int is expected,\n", + "# which triggers a Pydantic ValidationError during config parsing.\n", + "result_3 = run_failure_scenario(\n", + " client=client,\n", + " scenario_name=\"Invalid Training Hyperparameters\",\n", + " training_params={\n", + " \"model_path\": \"/opt/app-root/src\",\n", + " \"data_path\": \"/opt/app-root/src/data.jsonl\",\n", + " \"ckpt_output_dir\": \"/opt/app-root/src/checkpoints\",\n", + " \"data_output_dir\": \"/opt/app-root/src/sft-data\",\n", + " \"effective_batch_size\": 128,\n", + " \"num_epochs\": 1,\n", + " \"max_seq_len\": 4096,\n", + " \"max_batch_len\": \"not_a_number\",\n", + " \"learning_rate\": 5e-6,\n", + " },\n", + " expected_error=\"ValidationError\",\n", + " runtime=th_runtime,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cc9br1onzf", + "metadata": {}, + "outputs": [], + "source": [ + "# Scenario 4: OSFT — Invalid unfreeze_rank_ratio\n", + "# The OSFT algorithm validates unfreeze_rank_ratio with a manual range check\n", + "# (0.0 <= val <= 1.0), so passing a string triggers a TypeError.\n", + "result_4 = run_failure_scenario(\n", + " client=client,\n", + " scenario_name=\"OSFT Invalid Unfreeze Rank Ratio\",\n", + " training_params={\n", + " \"model_path\": \"/opt/app-root/src\",\n", + " \"data_path\": \"/opt/app-root/src/data.jsonl\",\n", + " \"ckpt_output_dir\": \"/opt/app-root/src/checkpoints\",\n", + " \"data_output_path\": \"/opt/app-root/src/osft-data\",\n", + " \"effective_batch_size\": 128,\n", + " \"num_epochs\": 1,\n", + " \"max_seq_len\": 2048,\n", + " \"max_tokens_per_gpu\": 32000,\n", + " \"learning_rate\": 5e-6,\n", + " \"unfreeze_rank_ratio\": \"not_a_number\", # Invalid: string instead of float\n", + " },\n", + " expected_error=\"'<=' not supported between instances of 'float' and 'str'\",\n", + " runtime=th_runtime,\n", + " algorithm=TrainingHubAlgorithms.OSFT,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "qj6u4whtjg", + "metadata": {}, + "outputs": [], + "source": [ + "# Scenario 5: LORA — Invalid model path\n", + "# LORA uses Unsloth's FastLanguageModel which raises a RuntimeError when\n", + "# the model path doesn't exist, before any LoRA-specific validation runs.\n", + "result_5 = run_failure_scenario(\n", + " client=client,\n", + " scenario_name=\"LORA Invalid Model Path\",\n", + " training_params={\n", + " \"model_path\": \"/nonexistent/model\",\n", + " \"data_path\": \"/opt/app-root/src/data.jsonl\",\n", + " \"ckpt_output_dir\": \"/opt/app-root/src/checkpoints\",\n", + " \"data_output_dir\": \"/opt/app-root/src/sft-data\",\n", + " \"micro_batch_size\": 16,\n", + " \"num_epochs\": 1,\n", + " \"max_seq_len\": 1024,\n", + " \"learning_rate\": 2e-4,\n", + " \"lora_r\": 16,\n", + " \"lora_alpha\": 32,\n", + " \"lora_dropout\": 0.0,\n", + " \"dataset_type\": \"chat_template\",\n", + " \"field_messages\": \"messages\",\n", + " },\n", + " expected_error=\"No config file found\",\n", + " runtime=th_runtime,\n", + " algorithm=TrainingHubAlgorithms.LORA_SFT,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-summary", + "metadata": {}, + "outputs": [], + "source": [ + "# Summary\n", + "results = {\n", + " \"Invalid Dataset Path\": result_1,\n", + " \"Invalid Model Path\": result_2,\n", + " \"Invalid Training Hyperparameters\": result_3,\n", + " \"OSFT Invalid Unfreeze Rank Ratio\": result_4,\n", + " \"LORA Invalid Model Path\": result_5,\n", + "}\n", + "\n", + "print(\"\\n\" + \"=\" * 60)\n", + "print(\"FAILURE SCENARIOS TEST SUMMARY\")\n", + "print(\"=\" * 60)\n", + "for name, passed in results.items():\n", + " status = \"PASSED\" if passed else \"FAILED\"\n", + " print(f\" {name}: {status}\")\n", + "\n", + "all_passed = all(results.values())\n", + "print(\"=\" * 60)\n", + "if all_passed:\n", + " print(\"NOTEBOOK_STATUS: SUCCESS\")\n", + "else:\n", + " failed = [name for name, passed in results.items() if not passed]\n", + " print(\"NOTEBOOK_STATUS: FAILURE\")\n", + " raise RuntimeError(f\"Failed scenarios: {', '.join(failed)}\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } }, - { - "cell_type": "code", - "id": "cell-k8s-setup", - "metadata": {}, - "outputs": [], - "source": "import os\nimport urllib3\nfrom kubernetes import client as k8s\n\n# Suppress InsecureRequestWarning since we use verify_ssl = False\nurllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)\n\napi_server = os.getenv(\"OPENSHIFT_API_URL\")\ntoken = os.getenv(\"NOTEBOOK_USER_TOKEN\")\nif not api_server or not token:\n raise RuntimeError(\"OPENSHIFT_API_URL and NOTEBOOK_USER_TOKEN environment variables are required\")\n\nconfiguration = k8s.Configuration()\nconfiguration.host = api_server\nconfiguration.verify_ssl = False\nconfiguration.api_key = {\"authorization\": f\"Bearer {token}\"}\napi_client = k8s.ApiClient(configuration)", - "execution_count": null - }, - { - "cell_type": "code", - "id": "cell-trainer-client", - "metadata": {}, - "outputs": [], - "source": [ - "from kubeflow.trainer import TrainerClient\n", - "from kubeflow.common.types import KubernetesBackendConfig\n", - "\n", - "backend_cfg = KubernetesBackendConfig(\n", - " client_configuration=api_client.configuration,\n", - ")\n", - "\n", - "client = TrainerClient(backend_cfg)\n", - "print(client)" - ], - "execution_count": null - }, - { - "cell_type": "code", - "id": "cell-find-runtime", - "metadata": {}, - "outputs": [], - "source": [ - "training_runtime_name = os.getenv(\"TRAINING_RUNTIME\")\n", - "if not training_runtime_name:\n", - " raise RuntimeError(\"TRAINING_RUNTIME environment variable is required\")\n", - "\n", - "th_runtime = None\n", - "for runtime in client.list_runtimes():\n", - " if runtime.name == training_runtime_name:\n", - " th_runtime = runtime\n", - " print(\"Found runtime: \" + str(th_runtime))\n", - " break\n", - "\n", - "if th_runtime is None:\n", - " raise RuntimeError(f\"Required runtime '{training_runtime_name}' not found\")" - ], - "execution_count": null - }, - { - "cell_type": "code", - "id": "cell-helper", - "metadata": {}, - "outputs": [], - "source": "import time\n\nfrom kubeflow.trainer.rhai import TrainingHubAlgorithms, TrainingHubTrainer\n\nNAMESPACE = os.getenv(\"NOTEBOOK_NAMESPACE\")\nprint(f\"Using namespace: {NAMESPACE}\")\n\n# Event reasons that indicate a container crash or failure\nFAILURE_EVENT_REASONS = {\"BackOff\", \"CrashLoopBackOff\", \"Failed\", \"OOMKilled\", \"OOMKilling\", \"Killing\"}\n\n\ndef check_job_failure(client, job_name):\n \"\"\"Check if a TrainJob has failed using SDK APIs (get_job and get_job_events).\n\n Returns (failure_confirmed, details) where failure is confirmed if:\n - get_job().status == \"Failed\", OR\n - get_job_events() contains crash-related events (BackOff, OOMKilled, etc.)\n \"\"\"\n failure_confirmed = False\n details = []\n\n # Check job status via get_job()\n try:\n job = client.get_job(name=job_name)\n if job.status == \"Failed\":\n failure_confirmed = True\n details.append(f\"job.status=Failed\")\n except Exception as e:\n print(f\" get_job() error: {e}\")\n\n # Check events via get_job_events()\n try:\n events = client.get_job_events(name=job_name)\n for event in events:\n reason = getattr(event, \"reason\", \"\") or \"\"\n if reason in FAILURE_EVENT_REASONS:\n failure_confirmed = True\n message = getattr(event, \"message\", \"\") or \"\"\n details.append(f\"event: reason={reason} message={message[:100]}\")\n except Exception as e:\n print(f\" get_job_events() error: {e}\")\n\n return failure_confirmed, details\n\n\ndef run_failure_scenario(client, scenario_name, training_params, expected_error, runtime, algorithm=TrainingHubAlgorithms.SFT):\n \"\"\"Submit a TrainJob expected to fail, verify the expected error in logs\n and that the job failure is detected via get_job()/get_job_events().\"\"\"\n print(f\"\\n{'='*60}\")\n print(f\"Scenario: {scenario_name}\")\n print(f\"Expected error: {expected_error}\")\n print(f\"{'='*60}\")\n\n job_name = None\n\n try:\n job_name = client.train(\n trainer=TrainingHubTrainer(\n algorithm=algorithm,\n func_args=training_params,\n resources_per_node={\n \"cpu\": 4,\n \"memory\": \"16Gi\",\n },\n ),\n runtime=runtime,\n )\n print(f\"TrainJob created: {job_name}\")\n\n # Wait for the job to start running\n try:\n client.wait_for_job_status(name=job_name, status={\"Running\", \"Failed\"}, timeout=300)\n except Exception as e:\n print(f\"Wait for Running/Failed raised: {e}\")\n\n # Poll until we see: (1) the expected error in logs, and (2) failure via SDK APIs\n error_found = False\n failure_confirmed = False\n deadline = time.time() + 300 # 5 minute timeout\n\n while time.time() < deadline:\n # Check logs for the expected error\n if not error_found:\n try:\n log_lines = list(client.get_job_logs(job_name, follow=False))\n log_text = \"\\n\".join(str(line) for line in log_lines)\n if expected_error in log_text:\n error_found = True\n print(f\"Found expected error in logs: '{expected_error}'\")\n except Exception as e:\n print(f\"Log fetch error (retrying): {e}\")\n\n # Check job failure via get_job() and get_job_events()\n if not failure_confirmed:\n failure_confirmed, details = check_job_failure(client, job_name)\n if failure_confirmed:\n print(f\"Job failure confirmed via SDK: {'; '.join(details)}\")\n\n if error_found and failure_confirmed:\n break\n\n time.sleep(15)\n\n # Verify results\n assert error_found, f\"Expected error '{expected_error}' not found in logs\"\n assert failure_confirmed, (\n f\"Job failure not confirmed via get_job()/get_job_events() \"\n f\"— final job status: {client.get_job(name=job_name).status}\"\n )\n\n print(f\"PASSED: {scenario_name}\")\n result = True\n\n except Exception as e:\n print(f\"FAILED: {scenario_name} - {e}\")\n result = False\n\n finally:\n # Cleanup\n if job_name:\n try:\n client.delete_job(job_name)\n print(f\"Deleted job: {job_name}\")\n except Exception as e:\n print(f\"Warning: failed to delete job {job_name}: {e}\")\n\n return result", - "execution_count": null - }, - { - "cell_type": "code", - "id": "cell-scenario-invalid-dataset", - "metadata": {}, - "outputs": [], - "source": "# Scenario 1: Invalid Dataset Path\nresult_1 = run_failure_scenario(\n client=client,\n scenario_name=\"Invalid Dataset Path\",\n training_params={\n \"model_path\": \"/opt/app-root/src\",\n \"data_path\": \"/nonexistent/dataset.jsonl\",\n \"ckpt_output_dir\": \"/opt/app-root/src/checkpoints\",\n \"data_output_dir\": \"/opt/app-root/src/sft-data\",\n \"effective_batch_size\": 128,\n \"num_epochs\": 1,\n \"max_seq_len\": 4096,\n \"max_batch_len\": 10000,\n \"learning_rate\": 5e-6,\n },\n expected_error=\"FileNotFoundError\",\n runtime=th_runtime,\n)", - "execution_count": null - }, - { - "cell_type": "code", - "id": "cell-scenario-invalid-model", - "metadata": {}, - "outputs": [], - "source": "# Scenario 2: Invalid Model Path\nresult_2 = run_failure_scenario(\n client=client,\n scenario_name=\"Invalid Model Path\",\n training_params={\n \"model_path\": \"/nonexistent/model\",\n \"data_path\": \"/opt/app-root/src/data.jsonl\",\n \"ckpt_output_dir\": \"/opt/app-root/src/checkpoints\",\n \"data_output_dir\": \"/opt/app-root/src/sft-data\",\n \"effective_batch_size\": 128,\n \"num_epochs\": 1,\n \"max_seq_len\": 4096,\n \"max_batch_len\": 10000,\n \"learning_rate\": 5e-6,\n },\n expected_error=\"FileNotFoundError\",\n runtime=th_runtime,\n)", - "execution_count": null - }, - { - "cell_type": "code", - "id": "cell-scenario-invalid-hyperparams", - "metadata": {}, - "outputs": [], - "source": "# Scenario 3: Invalid Training Hyperparameters\n# Uses a string value for max_batch_len where an int is expected,\n# which triggers a Pydantic ValidationError during config parsing.\nresult_3 = run_failure_scenario(\n client=client,\n scenario_name=\"Invalid Training Hyperparameters\",\n training_params={\n \"model_path\": \"/opt/app-root/src\",\n \"data_path\": \"/opt/app-root/src/data.jsonl\",\n \"ckpt_output_dir\": \"/opt/app-root/src/checkpoints\",\n \"data_output_dir\": \"/opt/app-root/src/sft-data\",\n \"effective_batch_size\": 128,\n \"num_epochs\": 1,\n \"max_seq_len\": 4096,\n \"max_batch_len\": \"not_a_number\",\n \"learning_rate\": 5e-6,\n },\n expected_error=\"ValidationError\",\n runtime=th_runtime,\n)", - "execution_count": null - }, - { - "cell_type": "code", - "id": "cc9br1onzf", - "source": "# Scenario 4: OSFT — Invalid unfreeze_rank_ratio\n# The OSFT algorithm validates unfreeze_rank_ratio with a manual range check\n# (0.0 <= val <= 1.0), so passing a string triggers a TypeError.\nresult_4 = run_failure_scenario(\n client=client,\n scenario_name=\"OSFT Invalid Unfreeze Rank Ratio\",\n training_params={\n \"model_path\": \"/opt/app-root/src\",\n \"data_path\": \"/opt/app-root/src/data.jsonl\",\n \"ckpt_output_dir\": \"/opt/app-root/src/checkpoints\",\n \"data_output_path\": \"/opt/app-root/src/osft-data\",\n \"effective_batch_size\": 128,\n \"num_epochs\": 1,\n \"max_seq_len\": 2048,\n \"max_tokens_per_gpu\": 32000,\n \"learning_rate\": 5e-6,\n \"unfreeze_rank_ratio\": \"not_a_number\", # Invalid: string instead of float\n },\n expected_error=\"'<=' not supported between instances of 'float' and 'str'\",\n runtime=th_runtime,\n algorithm=TrainingHubAlgorithms.OSFT,\n)", - "metadata": {}, - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "code", - "id": "qj6u4whtjg", - "source": "# Scenario 5: LORA — Invalid model path\n# LORA uses Unsloth's FastLanguageModel which raises a RuntimeError when\n# the model path doesn't exist, before any LoRA-specific validation runs.\nresult_5 = run_failure_scenario(\n client=client,\n scenario_name=\"LORA Invalid Model Path\",\n training_params={\n \"model_path\": \"/nonexistent/model\",\n \"data_path\": \"/opt/app-root/src/data.jsonl\",\n \"ckpt_output_dir\": \"/opt/app-root/src/checkpoints\",\n \"data_output_dir\": \"/opt/app-root/src/sft-data\",\n \"micro_batch_size\": 16,\n \"num_epochs\": 1,\n \"max_seq_len\": 1024,\n \"learning_rate\": 2e-4,\n \"lora_r\": 16,\n \"lora_alpha\": 32,\n \"lora_dropout\": 0.0,\n \"dataset_type\": \"chat_template\",\n \"field_messages\": \"messages\",\n },\n expected_error=\"No config file found\",\n runtime=th_runtime,\n algorithm=TrainingHubAlgorithms.LORA_SFT,\n)", - "metadata": {}, - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "code", - "id": "cell-summary", - "metadata": {}, - "outputs": [], - "source": "# Summary\nresults = {\n \"Invalid Dataset Path\": result_1,\n \"Invalid Model Path\": result_2,\n \"Invalid Training Hyperparameters\": result_3,\n \"OSFT Invalid Unfreeze Rank Ratio\": result_4,\n \"LORA Invalid Model Path\": result_5,\n}\n\nprint(\"\\n\" + \"=\" * 60)\nprint(\"FAILURE SCENARIOS TEST SUMMARY\")\nprint(\"=\" * 60)\nfor name, passed in results.items():\n status = \"PASSED\" if passed else \"FAILED\"\n print(f\" {name}: {status}\")\n\nall_passed = all(results.values())\nprint(\"=\" * 60)\nif all_passed:\n print(\"NOTEBOOK_STATUS: SUCCESS\")\nelse:\n failed = [name for name, passed in results.items() if not passed]\n print(\"NOTEBOOK_STATUS: FAILURE\")\n raise RuntimeError(f\"Failed scenarios: {', '.join(failed)}\")", - "execution_count": null - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "name": "python", - "version": "3.12.0" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} \ No newline at end of file + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tests/trainer/resources/lora.ipynb b/tests/trainer/resources/lora.ipynb index 48fc73a40..3901cbf60 100644 --- a/tests/trainer/resources/lora.ipynb +++ b/tests/trainer/resources/lora.ipynb @@ -1,627 +1,626 @@ { - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "%pip install datasets --quiet\n", - "# pip Install kubeflow SDK from main branch for testing\n", - "%pip install git+https://github.com/opendatahub-io/kubeflow-sdk.git@main --quiet" - ] + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%pip install datasets --quiet\n", + "# kubeflow SDK is installed by test harness via install_kubeflow.py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Standard library imports\n", + "import logging\n", + "import os\n", + "import sys\n", + "import time\n", + "from io import StringIO" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from dotenv import load_dotenv\n", + "load_dotenv()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from kubernetes import client as k8s, config as k8s_config\n", + "# Edit to match your specific settings\n", + "api_server = os.getenv(\"OPENSHIFT_API_URL\")\n", + "token = os.getenv(\"NOTEBOOK_USER_TOKEN\")\n", + "if not api_server or not token:\n", + " raise RuntimeError(\"OPENSHIFT_API_URL and NOTEBOOK_USER_TOKEN environment variables are required\")\n", + "PVC_NAME = os.getenv(\"SHARED_PVC_NAME\", \"\")\n", + "\n", + "if not PVC_NAME:\n", + " raise RuntimeError(\"SHARED_PVC_NAME environment variable is required\")\n", + "\n", + "configuration = k8s.Configuration()\n", + "configuration.host = api_server\n", + "# Un-comment if your cluster API server uses a self-signed certificate or an un-trusted CA\n", + "configuration.verify_ssl = False\n", + "configuration.api_key = {\"authorization\": f\"Bearer {token}\"}\n", + "api_client = k8s.ApiClient(configuration)\n", + "\n", + "PVC_MOUNT_PATH = \"/opt/app-root/src\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Converting the format of the intial messages.\n", + "def convert_to_messages(example):\n", + " \"\"\"\n", + " Convert a sql-create-context example to chat template format.\n", + " \n", + " The user provides the database schema and question.\n", + " The assistant responds with the SQL query.\n", + " \"\"\"\n", + " user_message = f\"\"\"Given the following database schema:\n", + "\n", + "{example['context']}\n", + "\n", + "Write a SQL query to answer this question: {example['question']}\"\"\"\n", + " \n", + " assistant_message = example['answer']\n", + " \n", + " return {\n", + " \"messages\": [\n", + " {\"role\": \"user\", \"content\": user_message},\n", + " {\"role\": \"assistant\", \"content\": assistant_message}\n", + " ]\n", + " }" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import gzip\n", + "import shutil\n", + "import socket\n", + "\n", + "import boto3\n", + "from botocore.config import Config as BotoConfig\n", + "from botocore.exceptions import ClientError\n", + "\n", + "# --- Global networking safety net: cap all socket operations ---\n", + "socket.setdefaulttimeout(10) # seconds\n", + "\n", + "# Notebook's PVC mount path (per Notebook CR). Training pods will mount the same PVC at /opt/app-root/src\n", + "PVC_NOTEBOOK_PATH = \"/opt/app-root/src/\"\n", + "DATASET_ROOT_NOTEBOOK = PVC_NOTEBOOK_PATH\n", + "TXT_SQL_DIR = os.path.join(DATASET_ROOT_NOTEBOOK, \"txt-sql-data\", \"train\")\n", + "MODEL_DIR = os.path.join(DATASET_ROOT_NOTEBOOK, \"Qwen\", \"Qwen2.5-1.5B-Instruct\")\n", + "os.makedirs(TXT_SQL_DIR, exist_ok=True)\n", + "os.makedirs(MODEL_DIR, exist_ok=True)\n", + "\n", + "# Env config for S3/MinIO\n", + "s3_endpoint = os.getenv(\"AWS_DEFAULT_ENDPOINT\", \"\")\n", + "s3_access_key = os.getenv(\"AWS_ACCESS_KEY_ID\", \"\")\n", + "s3_secret_key = os.getenv(\"AWS_SECRET_ACCESS_KEY\", \"\")\n", + "s3_bucket = os.getenv(\"AWS_STORAGE_BUCKET\", \"\")\n", + "s3_prefix = os.getenv(\"AWS_STORAGE_BUCKET_LORA_DIR\", \"\") \n", + "\n", + "data_download_successful = False\n", + "\n", + "def stream_download(s3, bucket, key, dst):\n", + " \"\"\"\n", + " Download an object from S3/MinIO using get_object and streaming reads.\n", + " Returns True on success, False on any error.\n", + " \"\"\"\n", + " print(f\"[notebook] STREAM download s3://{bucket}/{key} -> {dst}\")\n", + " t0 = time.time()\n", + "\n", + " try:\n", + " resp = s3.get_object(Bucket=bucket, Key=key)\n", + " except ClientError as e:\n", + " err = e.response.get(\"Error\", {})\n", + " print(f\"[notebook] CLIENT ERROR (get_object) for {key}: {err}\")\n", + " return False\n", + " except Exception as e:\n", + " print(f\"[notebook] OTHER ERROR (get_object) for {key}: {e}\")\n", + " return False\n", + "\n", + " body = resp[\"Body\"]\n", + " try:\n", + " with open(dst, \"wb\") as f:\n", + " while True:\n", + " try:\n", + " chunk = body.read(1024 * 1024) # 1MB per chunk\n", + " except socket.timeout as e:\n", + " print(f\"[notebook] socket.timeout while reading {key}: {e}\")\n", + " return False\n", + " if not chunk:\n", + " break\n", + " f.write(chunk)\n", + " except Exception as e:\n", + " print(f\"[notebook] ERROR writing to {dst} for {key}: {e}\")\n", + " return False\n", + "\n", + " t1 = time.time()\n", + " print(f\"[notebook] DONE stream {key} in {t1 - t0:.2f}s\")\n", + " return True\n", + "\n", + "\n", + "if s3_endpoint and s3_bucket:\n", + " try:\n", + " # Normalize endpoint URL\n", + " endpoint_url = (\n", + " s3_endpoint\n", + " if s3_endpoint.startswith(\"http\")\n", + " else f\"https://{s3_endpoint}\"\n", + " )\n", + " prefix = (s3_prefix or \"\").strip(\"/\")\n", + "\n", + " print(\n", + " f\"[notebook] S3 configured: \"\n", + " f\"endpoint={endpoint_url}, bucket={s3_bucket}, prefix={prefix or ''}\"\n", + " )\n", + "\n", + " # Boto config: single attempt, reasonable connect/read timeouts\n", + " boto_cfg = BotoConfig(\n", + " signature_version=\"s3v4\",\n", + " s3={\"addressing_style\": \"path\"},\n", + " retries={\"max_attempts\": 1, \"mode\": \"standard\"},\n", + " connect_timeout=5,\n", + " read_timeout=10,\n", + " )\n", + "\n", + " # Create S3/MinIO client\n", + " s3 = boto3.client(\n", + " \"s3\",\n", + " endpoint_url=endpoint_url,\n", + " aws_access_key_id=s3_access_key,\n", + " aws_secret_access_key=s3_secret_key,\n", + " config=boto_cfg,\n", + " verify=False,\n", + " )\n", + "\n", + " # List and download all objects under the prefix\n", + " paginator = s3.get_paginator(\"list_objects_v2\")\n", + " pulled_any = False\n", + " file_count = 0\n", + " \n", + " print(f\"[notebook] Starting S3 download from prefix: {prefix}\")\n", + " for page in paginator.paginate(Bucket=s3_bucket, Prefix=prefix or \"\"):\n", + " contents = page.get(\"Contents\", [])\n", + " if not contents:\n", + " print(f\"[notebook] No contents found in this page\")\n", + " continue\n", + " \n", + " print(f\"[notebook] Found {len(contents)} objects in this page\")\n", + "\n", + " for obj in contents:\n", + " key = obj[\"Key\"]\n", + " file_count += 1\n", + "\n", + " # Skip \"directory markers\"\n", + " if key.endswith(\"/\"):\n", + " print(f\"[notebook] Skipping directory marker: {key}\")\n", + " continue\n", + "\n", + " # Determine relative path under prefix for local storage\n", + " rel = key[len(prefix):].lstrip(\"/\") if prefix else key\n", + " print(f\"[notebook] Processing key={key}, rel={rel}\")\n", + " \n", + " # Route to appropriate directory based on content type\n", + " if \"table-gpt\" in rel.lower() or rel.endswith(\".jsonl\"):\n", + " dst = os.path.join(TXT_SQL_DIR, os.path.basename(rel))\n", + " print(f\"[notebook] Routing to dataset dir: {dst}\")\n", + " elif \"qwen\" in rel.lower() or any(rel.endswith(ext) for ext in [\".bin\", \".json\", \".model\", \".safetensors\", \".txt\"]):\n", + " # Preserve directory structure for model files\n", + " dst = os.path.join(MODEL_DIR, rel.split(\"Qwen2.5-1.5B-Instruct/\")[-1] if \"Qwen2.5-1.5B-Instruct\" in rel else os.path.basename(rel))\n", + " print(f\"[notebook] Routing to model dir: {dst}\")\n", + " else:\n", + " # Default: use the relative path as-is\n", + " dst = os.path.join(DATASET_ROOT_NOTEBOOK, rel)\n", + " print(f\"[notebook] Routing to default dir: {dst}\")\n", + " \n", + " os.makedirs(os.path.dirname(dst), exist_ok=True)\n", + "\n", + " # Download only if missing\n", + " if not os.path.exists(dst):\n", + " ok = stream_download(s3, s3_bucket, key, dst)\n", + " if not ok:\n", + " print(f\"[notebook] Download failed for {key}\")\n", + " continue\n", + " pulled_any = True\n", + " else:\n", + " print(f\"[notebook] Skipping existing file {dst}\")\n", + " pulled_any = True\n", + "\n", + " # If the file is .gz, decompress and remove the .gz\n", + " if dst.endswith(\".gz\") and os.path.exists(dst):\n", + " out_path = os.path.splitext(dst)[0]\n", + " if not os.path.exists(out_path):\n", + " print(f\"[notebook] Decompressing {dst} -> {out_path}\")\n", + " try:\n", + " with gzip.open(dst, \"rb\") as f_in, open(out_path, \"wb\") as f_out:\n", + " shutil.copyfileobj(f_in, f_out)\n", + " except Exception as e:\n", + " print(f\"[notebook] Failed to decompress {dst}: {e}\")\n", + " else:\n", + " try:\n", + " os.remove(dst)\n", + " except Exception:\n", + " pass\n", + "\n", + " if pulled_any:\n", + " print(f\"[notebook] ✓ S3 download successful. Processed {file_count} files\")\n", + " data_download_successful = True\n", + " else:\n", + " print(f\"[notebook] ✗ S3 download found no files to download\")\n", + "\n", + " except Exception as e:\n", + " print(f\"[notebook] ✗ S3 fetch failed: {e}\")\n", + " import traceback\n", + " traceback.print_exc()\n", + " print(\"[notebook] Will attempt HuggingFace fallback...\")\n", + "else:\n", + " print(\"[notebook] S3 not configured (missing endpoint or bucket env vars)\")\n", + "\n", + "# Fallback to HuggingFace if S3 was not configured or failed (requires internet)\n", + "if not data_download_successful:\n", + " print(\"[notebook] Attempting HuggingFace dataset download (requires internet)...\")\n", + " try:\n", + " import json\n", + " import random\n", + " from datasets import load_dataset\n", + "\n", + " # Load the Table-GPT dataset\n", + " print(\"[notebook] Loading Table-GPT dataset from HuggingFace...\")\n", + " # Load the dataset\n", + " dataset = load_dataset(\"b-mc2/sql-create-context\", split=\"train\")\n", + "\n", + " TRAIN_SIZE = 100 # Adjust based on your time/compute budget\n", + "\n", + " # Shuffle and select a subset\n", + " train_dataset = dataset.shuffle(seed=42).select(range(min(TRAIN_SIZE, len(dataset))))\n", + "\n", + " # Convert to messages format\n", + " train_data = [convert_to_messages(example) for example in train_dataset]\n", + "\n", + " # Save the subset to a JSONL file\n", + " output_file = os.path.join(TXT_SQL_DIR, \"train_All_100.jsonl\")\n", + " with open(output_file, \"w\") as f:\n", + " for example in train_data:\n", + " f.write(json.dumps(example) + \"\\n\")\n", + "\n", + " print(f\"[notebook] ✓ HuggingFace download successful. Subset saved to {output_file}\")\n", + " data_download_successful = True\n", + "\n", + " except Exception as hf_error:\n", + " print(f\"[notebook] ✗ HuggingFace download failed: {hf_error}\")\n", + " import traceback\n", + " traceback.print_exc()\n", + " raise RuntimeError(\n", + " \"Failed to download dataset from both S3 and HuggingFace. \"\n", + " \"In disconnected environments, ensure S3/MinIO is configured with the required data. \"\n", + " \"In connected environments, check your internet connection and credentials.\"\n", + " ) from hf_error\n", + "\n", + "# Verify dataset file exists\n", + "dataset_file = os.path.join(TXT_SQL_DIR, \"train_All_100.jsonl\")\n", + "if os.path.exists(dataset_file):\n", + " print(f\"[notebook] ✓ Dataset ready: {dataset_file}\")\n", + "else:\n", + " raise RuntimeError(f\"Dataset file not found: {dataset_file}\")\n", + "\n", + "# Verify model directory has files (model will be downloaded during training if not present)\n", + "if os.path.exists(MODEL_DIR) and os.listdir(MODEL_DIR):\n", + " print(f\"[notebook] ✓ Model files ready in: {MODEL_DIR}\")\n", + " print(f\"[notebook] Model files: {os.listdir(MODEL_DIR)[:5]}...\") # Show first 5 files\n", + "else:\n", + " print(f\"[notebook] Note: Model directory is empty: {MODEL_DIR}\")\n", + " print(\"[notebook] Training will download model from HuggingFace during execution\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Model download - use S3 if available, otherwise HuggingFace\n", + "if os.path.exists(MODEL_DIR) and os.listdir(MODEL_DIR):\n", + " model_path = MODEL_DIR\n", + " print(f\"✓ Using local model from S3: {model_path}\")\n", + "else:\n", + " # Download from HuggingFace\n", + " print(\"[notebook] Model not found in S3, downloading from HuggingFace...\")\n", + " from huggingface_hub import snapshot_download\n", + " \n", + " token = os.getenv(\"HUGGINGFACE_HUB_TOKEN\")\n", + " model_path = snapshot_download(\n", + " repo_id=\"Qwen/Qwen2.5-1.5B-Instruct\",\n", + " local_dir=MODEL_DIR,\n", + " token=token,\n", + " resume_download=True,\n", + " local_dir_use_symlinks=False,\n", + " )\n", + " print(f\"✓ Model downloaded to: {model_path}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Training configuration\n", + "MODEL_NAME = \"Qwen/Qwen2.5-1.5B-Instruct\"\n", + "\n", + "# You can also try these alternatives:\n", + "# MODEL_NAME = \"Qwen/Qwen2.5-3B-Instruct\" # Larger, more capable\n", + "# MODEL_NAME = \"Qwen/Qwen2.5-0.5B-Instruct\" # Smaller, faster training\n", + "# MODEL_NAME = \"meta-llama/Llama-3.2-1B-Instruct\" # Alternative architecture\n", + "\n", + "# LoRA configuration\n", + "LORA_R = 16 # Rank - start small, increase if needed\n", + "LORA_ALPHA = 32 # Alpha - typically 2x rank\n", + "LORA_DROPOUT = 0.0 # Dropout - 0.0 is optimized for Unsloth\n", + "\n", + "# Training configuration\n", + "NUM_EPOCHS = 2 # More epochs = better learning, longer training\n", + "LEARNING_RATE = 2e-4 # Standard LoRA learning rate\n", + "MAX_SEQ_LEN = 1024 # Maximum sequence length\n", + "MICRO_BATCH_SIZE = 16 # Batch size per GPU (reduce if OOM)\n", + "GRADIENT_ACCUMULATION = 4 # Effective batch = micro_batch * grad_accum\n", + "\n", + "# QLoRA settings (set to True to enable 4-bit quantization)\n", + "USE_QLORA = True # Set to True if you have limited GPU memory\n", + "\n", + "params = {\n", + " # Model and data path\n", + " 'model_path': MODEL_NAME,\n", + " 'data_path': \"/opt/app-root/src/txt-sql-data/train/train_All_100.jsonl\",\n", + " 'ckpt_output_dir': \"/opt/app-root/src/checkpoints-logs-dir\",\n", + " 'data_output_path': \"/opt/app-root/src/lora-json/_data\",\n", + " # Important for LORA\n", + " 'lr_scheduler': \"cosine\",\n", + " 'warmup_steps': 0,\n", + " 'seed': 42,\n", + " # LoRA configuration\n", + " 'lora_r': LORA_R,\n", + " 'lora_alpha': LORA_ALPHA,\n", + " 'lora_dropout': LORA_DROPOUT,\n", + "\n", + " # Training configuration\n", + " 'num_epochs': NUM_EPOCHS,\n", + " 'learning_rate': LEARNING_RATE,\n", + " 'micro_batch_size': MICRO_BATCH_SIZE,\n", + " 'max_seq_len': MAX_SEQ_LEN,\n", + " 'gradient_accumulation_steps': GRADIENT_ACCUMULATION,\n", + "\n", + " # Dataset format\n", + " 'dataset_type' : \"chat_template\",\n", + " 'field_messages' : \"messages\",\n", + " # Quantization\n", + " 'load_in_4bit': USE_QLORA,\n", + " #GPU configuration\n", + " 'nproc_per_node' : 2,\n", + " 'nnodes' : 2,\n", + " # Logging\n", + " 'logging_steps': 10,\n", + " 'save_steps': 200,\n", + " 'save_total_limit': 3,\n", + "\n", + " # Model Checkpointing\n", + " 'save_final_checkpoint': True,\n", + " 'checkpoint_at_epoch': 2,\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from kubeflow.trainer import TrainerClient\n", + "from kubeflow.trainer.rhai import TrainingHubAlgorithms\n", + "from kubeflow.trainer.rhai import TrainingHubTrainer\n", + "from kubeflow.common.types import KubernetesBackendConfig\n", + "\n", + "backend_cfg = KubernetesBackendConfig(client_configuration=api_client.configuration)\n", + "client = TrainerClient(backend_cfg)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "training_runtime_name = os.getenv(\"TRAINING_RUNTIME\")\n", + "\n", + "if not training_runtime_name:\n", + " raise RuntimeError(\"TRAINING_RUNTIME environment variable is required\")\n", + "\n", + "th_runtime = None\n", + "for runtime in client.list_runtimes():\n", + " if runtime.name == training_runtime_name:\n", + " th_runtime = runtime\n", + " print(\"Found runtime: \" + str(th_runtime))\n", + " break\n", + "\n", + "if th_runtime is None:\n", + " raise RuntimeError(f\"Required runtime '{training_runtime_name}' not found\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from kubeflow.trainer.options.kubernetes import (\n", + " PodTemplateOverrides,\n", + " PodTemplateOverride,\n", + " PodSpecOverride,\n", + " ContainerOverride,\n", + ")\n", + "\n", + "cache_root = \"/opt/app-root/src/.cache/huggingface\"\n", + "triton_cache = \"/opt/app-root/src/.triton\"\n", + "\n", + "job_name = client.train(\n", + " trainer=TrainingHubTrainer(\n", + " algorithm=TrainingHubAlgorithms.LORA_SFT,\n", + " func_args=params,\n", + " env={\n", + " \"HF_HOME\": cache_root,\n", + " \"TRITON_CACHE_DIR\": triton_cache,\n", + " \"XDG_CACHE_HOME\": \"/opt/app-root/src/.cache\",\n", + " \"NCCL_DEBUG\": \"INFO\",\n", + " },\n", + " resources_per_node={\n", + " \"cpu\": 4,\n", + " \"memory\": \"32Gi\",\n", + " \"nvidia.com/gpu\": 1\n", + " },\n", + " ),\n", + " options=[\n", + " PodTemplateOverrides(\n", + " PodTemplateOverride(\n", + " target_jobs=[\"node\"],\n", + " spec=PodSpecOverride(\n", + " volumes=[\n", + " {\"name\": \"work\", \"persistentVolumeClaim\": {\"claimName\": PVC_NAME}},\n", + " ],\n", + " containers=[\n", + " ContainerOverride(\n", + " name=\"node\",\n", + " volume_mounts=[\n", + " {\"name\": \"work\", \"mountPath\": \"/opt/app-root/src\", \"readOnly\": False},\n", + " ],\n", + " )\n", + " ],\n", + " ),\n", + " )\n", + " )\n", + " ],\n", + " runtime=th_runtime,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Wait for the running status, then wait for completion or failure\n", + "# Using reasonable timeout for LORA training\n", + "client.wait_for_job_status(name=job_name, status={\"Running\"}, timeout=300)\n", + "client.wait_for_job_status(name=job_name, status={\"Complete\", \"Failed\"}, timeout=1800) # 30 minutes for training\n", + "\n", + "# Get job details and logs\n", + "job = client.get_job(name=job_name)\n", + "pod_logs = client.get_job_logs(name=job_name, follow=False)\n", + "\n", + "# Collect all log lines from the generator into a list\n", + "logs = list(pod_logs)\n", + "log_text = \"\\n\".join(str(line) for line in logs)\n", + "\n", + "print(f\"Training job final status: {job.status}\")\n", + "\n", + "# Check 1: Job status must not be \"Failed\" \n", + "if job.status == \"Failed\":\n", + " print(f\"ERROR: Training job '{job_name}' has Failed status\")\n", + " print(\"Last 30 lines of logs:\")\n", + " for line in logs[-30:]:\n", + " print(line)\n", + " raise RuntimeError(f\"Training job '{job_name}' failed\")\n", + "\n", + "# Check 2: Look for the training completion message in logs\n", + "# This is critical because the training script may catch exceptions and exit 0\n", + "if \"[PY] LORA_SFT training complete. Result=\" not in log_text:\n", + " print(f\"ERROR: Training completion message not found in logs\")\n", + " print(\"Last 50 lines of logs:\")\n", + " for line in logs[-50:]:\n", + " print(line)\n", + " raise RuntimeError(f\"Training did not complete successfully - missing completion message\")\n", + "\n", + "print(f\"✓ Training job '{job_name}' completed successfully\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "for c in client.get_job(name=job_name).steps:\n", + " print(f\"Step: {c.name}, Status: {c.status}, Devices: {c.device} x {c.device_count}\\n\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "logs = client.get_job_logs(name=job_name, follow=False)\n", + "\n", + "# Collect all log lines from the generator into a list\n", + "logs = list(pod_logs)\n", + "log_text = \"\\n\".join(str(line) for line in logs)\n", + "print(log_text)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "client.delete_job(job_name)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.14.1" + } }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Standard library imports\n", - "import logging\n", - "import os\n", - "import sys\n", - "import time\n", - "from io import StringIO" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from dotenv import load_dotenv\n", - "load_dotenv()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from kubernetes import client as k8s, config as k8s_config\n", - "# Edit to match your specific settings\n", - "api_server = os.getenv(\"OPENSHIFT_API_URL\")\n", - "token = os.getenv(\"NOTEBOOK_USER_TOKEN\")\n", - "if not api_server or not token:\n", - " raise RuntimeError(\"OPENSHIFT_API_URL and NOTEBOOK_USER_TOKEN environment variables are required\")\n", - "PVC_NAME = os.getenv(\"SHARED_PVC_NAME\", \"\")\n", - "\n", - "if not PVC_NAME:\n", - " raise RuntimeError(\"SHARED_PVC_NAME environment variable is required\")\n", - "\n", - "configuration = k8s.Configuration()\n", - "configuration.host = api_server\n", - "# Un-comment if your cluster API server uses a self-signed certificate or an un-trusted CA\n", - "configuration.verify_ssl = False\n", - "configuration.api_key = {\"authorization\": f\"Bearer {token}\"}\n", - "api_client = k8s.ApiClient(configuration)\n", - "\n", - "PVC_MOUNT_PATH = \"/opt/app-root/src\"" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Converting the format of the intial messages.\n", - "def convert_to_messages(example):\n", - " \"\"\"\n", - " Convert a sql-create-context example to chat template format.\n", - " \n", - " The user provides the database schema and question.\n", - " The assistant responds with the SQL query.\n", - " \"\"\"\n", - " user_message = f\"\"\"Given the following database schema:\n", - "\n", - "{example['context']}\n", - "\n", - "Write a SQL query to answer this question: {example['question']}\"\"\"\n", - " \n", - " assistant_message = example['answer']\n", - " \n", - " return {\n", - " \"messages\": [\n", - " {\"role\": \"user\", \"content\": user_message},\n", - " {\"role\": \"assistant\", \"content\": assistant_message}\n", - " ]\n", - " }" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", - "import gzip\n", - "import shutil\n", - "import socket\n", - "\n", - "import boto3\n", - "from botocore.config import Config as BotoConfig\n", - "from botocore.exceptions import ClientError\n", - "\n", - "# --- Global networking safety net: cap all socket operations ---\n", - "socket.setdefaulttimeout(10) # seconds\n", - "\n", - "# Notebook's PVC mount path (per Notebook CR). Training pods will mount the same PVC at /opt/app-root/src\n", - "PVC_NOTEBOOK_PATH = \"/opt/app-root/src/\"\n", - "DATASET_ROOT_NOTEBOOK = PVC_NOTEBOOK_PATH\n", - "TXT_SQL_DIR = os.path.join(DATASET_ROOT_NOTEBOOK, \"txt-sql-data\", \"train\")\n", - "MODEL_DIR = os.path.join(DATASET_ROOT_NOTEBOOK, \"Qwen\", \"Qwen2.5-1.5B-Instruct\")\n", - "os.makedirs(TXT_SQL_DIR, exist_ok=True)\n", - "os.makedirs(MODEL_DIR, exist_ok=True)\n", - "\n", - "# Env config for S3/MinIO\n", - "s3_endpoint = os.getenv(\"AWS_DEFAULT_ENDPOINT\", \"\")\n", - "s3_access_key = os.getenv(\"AWS_ACCESS_KEY_ID\", \"\")\n", - "s3_secret_key = os.getenv(\"AWS_SECRET_ACCESS_KEY\", \"\")\n", - "s3_bucket = os.getenv(\"AWS_STORAGE_BUCKET\", \"\")\n", - "s3_prefix = os.getenv(\"AWS_STORAGE_BUCKET_LORA_DIR\", \"\") \n", - "\n", - "data_download_successful = False\n", - "\n", - "def stream_download(s3, bucket, key, dst):\n", - " \"\"\"\n", - " Download an object from S3/MinIO using get_object and streaming reads.\n", - " Returns True on success, False on any error.\n", - " \"\"\"\n", - " print(f\"[notebook] STREAM download s3://{bucket}/{key} -> {dst}\")\n", - " t0 = time.time()\n", - "\n", - " try:\n", - " resp = s3.get_object(Bucket=bucket, Key=key)\n", - " except ClientError as e:\n", - " err = e.response.get(\"Error\", {})\n", - " print(f\"[notebook] CLIENT ERROR (get_object) for {key}: {err}\")\n", - " return False\n", - " except Exception as e:\n", - " print(f\"[notebook] OTHER ERROR (get_object) for {key}: {e}\")\n", - " return False\n", - "\n", - " body = resp[\"Body\"]\n", - " try:\n", - " with open(dst, \"wb\") as f:\n", - " while True:\n", - " try:\n", - " chunk = body.read(1024 * 1024) # 1MB per chunk\n", - " except socket.timeout as e:\n", - " print(f\"[notebook] socket.timeout while reading {key}: {e}\")\n", - " return False\n", - " if not chunk:\n", - " break\n", - " f.write(chunk)\n", - " except Exception as e:\n", - " print(f\"[notebook] ERROR writing to {dst} for {key}: {e}\")\n", - " return False\n", - "\n", - " t1 = time.time()\n", - " print(f\"[notebook] DONE stream {key} in {t1 - t0:.2f}s\")\n", - " return True\n", - "\n", - "\n", - "if s3_endpoint and s3_bucket:\n", - " try:\n", - " # Normalize endpoint URL\n", - " endpoint_url = (\n", - " s3_endpoint\n", - " if s3_endpoint.startswith(\"http\")\n", - " else f\"https://{s3_endpoint}\"\n", - " )\n", - " prefix = (s3_prefix or \"\").strip(\"/\")\n", - "\n", - " print(\n", - " f\"[notebook] S3 configured: \"\n", - " f\"endpoint={endpoint_url}, bucket={s3_bucket}, prefix={prefix or ''}\"\n", - " )\n", - "\n", - " # Boto config: single attempt, reasonable connect/read timeouts\n", - " boto_cfg = BotoConfig(\n", - " signature_version=\"s3v4\",\n", - " s3={\"addressing_style\": \"path\"},\n", - " retries={\"max_attempts\": 1, \"mode\": \"standard\"},\n", - " connect_timeout=5,\n", - " read_timeout=10,\n", - " )\n", - "\n", - " # Create S3/MinIO client\n", - " s3 = boto3.client(\n", - " \"s3\",\n", - " endpoint_url=endpoint_url,\n", - " aws_access_key_id=s3_access_key,\n", - " aws_secret_access_key=s3_secret_key,\n", - " config=boto_cfg,\n", - " verify=False,\n", - " )\n", - "\n", - " # List and download all objects under the prefix\n", - " paginator = s3.get_paginator(\"list_objects_v2\")\n", - " pulled_any = False\n", - " file_count = 0\n", - " \n", - " print(f\"[notebook] Starting S3 download from prefix: {prefix}\")\n", - " for page in paginator.paginate(Bucket=s3_bucket, Prefix=prefix or \"\"):\n", - " contents = page.get(\"Contents\", [])\n", - " if not contents:\n", - " print(f\"[notebook] No contents found in this page\")\n", - " continue\n", - " \n", - " print(f\"[notebook] Found {len(contents)} objects in this page\")\n", - "\n", - " for obj in contents:\n", - " key = obj[\"Key\"]\n", - " file_count += 1\n", - "\n", - " # Skip \"directory markers\"\n", - " if key.endswith(\"/\"):\n", - " print(f\"[notebook] Skipping directory marker: {key}\")\n", - " continue\n", - "\n", - " # Determine relative path under prefix for local storage\n", - " rel = key[len(prefix):].lstrip(\"/\") if prefix else key\n", - " print(f\"[notebook] Processing key={key}, rel={rel}\")\n", - " \n", - " # Route to appropriate directory based on content type\n", - " if \"table-gpt\" in rel.lower() or rel.endswith(\".jsonl\"):\n", - " dst = os.path.join(TXT_SQL_DIR, os.path.basename(rel))\n", - " print(f\"[notebook] Routing to dataset dir: {dst}\")\n", - " elif \"qwen\" in rel.lower() or any(rel.endswith(ext) for ext in [\".bin\", \".json\", \".model\", \".safetensors\", \".txt\"]):\n", - " # Preserve directory structure for model files\n", - " dst = os.path.join(MODEL_DIR, rel.split(\"Qwen2.5-1.5B-Instruct/\")[-1] if \"Qwen2.5-1.5B-Instruct\" in rel else os.path.basename(rel))\n", - " print(f\"[notebook] Routing to model dir: {dst}\")\n", - " else:\n", - " # Default: use the relative path as-is\n", - " dst = os.path.join(DATASET_ROOT_NOTEBOOK, rel)\n", - " print(f\"[notebook] Routing to default dir: {dst}\")\n", - " \n", - " os.makedirs(os.path.dirname(dst), exist_ok=True)\n", - "\n", - " # Download only if missing\n", - " if not os.path.exists(dst):\n", - " ok = stream_download(s3, s3_bucket, key, dst)\n", - " if not ok:\n", - " print(f\"[notebook] Download failed for {key}\")\n", - " continue\n", - " pulled_any = True\n", - " else:\n", - " print(f\"[notebook] Skipping existing file {dst}\")\n", - " pulled_any = True\n", - "\n", - " # If the file is .gz, decompress and remove the .gz\n", - " if dst.endswith(\".gz\") and os.path.exists(dst):\n", - " out_path = os.path.splitext(dst)[0]\n", - " if not os.path.exists(out_path):\n", - " print(f\"[notebook] Decompressing {dst} -> {out_path}\")\n", - " try:\n", - " with gzip.open(dst, \"rb\") as f_in, open(out_path, \"wb\") as f_out:\n", - " shutil.copyfileobj(f_in, f_out)\n", - " except Exception as e:\n", - " print(f\"[notebook] Failed to decompress {dst}: {e}\")\n", - " else:\n", - " try:\n", - " os.remove(dst)\n", - " except Exception:\n", - " pass\n", - "\n", - " if pulled_any:\n", - " print(f\"[notebook] ✓ S3 download successful. Processed {file_count} files\")\n", - " data_download_successful = True\n", - " else:\n", - " print(f\"[notebook] ✗ S3 download found no files to download\")\n", - "\n", - " except Exception as e:\n", - " print(f\"[notebook] ✗ S3 fetch failed: {e}\")\n", - " import traceback\n", - " traceback.print_exc()\n", - " print(\"[notebook] Will attempt HuggingFace fallback...\")\n", - "else:\n", - " print(\"[notebook] S3 not configured (missing endpoint or bucket env vars)\")\n", - "\n", - "# Fallback to HuggingFace if S3 was not configured or failed (requires internet)\n", - "if not data_download_successful:\n", - " print(\"[notebook] Attempting HuggingFace dataset download (requires internet)...\")\n", - " try:\n", - " import json\n", - " import random\n", - " from datasets import load_dataset\n", - "\n", - " # Load the Table-GPT dataset\n", - " print(\"[notebook] Loading Table-GPT dataset from HuggingFace...\")\n", - " # Load the dataset\n", - " dataset = load_dataset(\"b-mc2/sql-create-context\", split=\"train\")\n", - "\n", - " TRAIN_SIZE = 100 # Adjust based on your time/compute budget\n", - "\n", - " # Shuffle and select a subset\n", - " train_dataset = dataset.shuffle(seed=42).select(range(min(TRAIN_SIZE, len(dataset))))\n", - "\n", - " # Convert to messages format\n", - " train_data = [convert_to_messages(example) for example in train_dataset]\n", - "\n", - " # Save the subset to a JSONL file\n", - " output_file = os.path.join(TXT_SQL_DIR, \"train_All_100.jsonl\")\n", - " with open(output_file, \"w\") as f:\n", - " for example in train_data:\n", - " f.write(json.dumps(example) + \"\\n\")\n", - "\n", - " print(f\"[notebook] ✓ HuggingFace download successful. Subset saved to {output_file}\")\n", - " data_download_successful = True\n", - "\n", - " except Exception as hf_error:\n", - " print(f\"[notebook] ✗ HuggingFace download failed: {hf_error}\")\n", - " import traceback\n", - " traceback.print_exc()\n", - " raise RuntimeError(\n", - " \"Failed to download dataset from both S3 and HuggingFace. \"\n", - " \"In disconnected environments, ensure S3/MinIO is configured with the required data. \"\n", - " \"In connected environments, check your internet connection and credentials.\"\n", - " ) from hf_error\n", - "\n", - "# Verify dataset file exists\n", - "dataset_file = os.path.join(TXT_SQL_DIR, \"train_All_100.jsonl\")\n", - "if os.path.exists(dataset_file):\n", - " print(f\"[notebook] ✓ Dataset ready: {dataset_file}\")\n", - "else:\n", - " raise RuntimeError(f\"Dataset file not found: {dataset_file}\")\n", - "\n", - "# Verify model directory has files (model will be downloaded during training if not present)\n", - "if os.path.exists(MODEL_DIR) and os.listdir(MODEL_DIR):\n", - " print(f\"[notebook] ✓ Model files ready in: {MODEL_DIR}\")\n", - " print(f\"[notebook] Model files: {os.listdir(MODEL_DIR)[:5]}...\") # Show first 5 files\n", - "else:\n", - " print(f\"[notebook] Note: Model directory is empty: {MODEL_DIR}\")\n", - " print(\"[notebook] Training will download model from HuggingFace during execution\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Model download - use S3 if available, otherwise HuggingFace\n", - "if os.path.exists(MODEL_DIR) and os.listdir(MODEL_DIR):\n", - " model_path = MODEL_DIR\n", - " print(f\"✓ Using local model from S3: {model_path}\")\n", - "else:\n", - " # Download from HuggingFace\n", - " print(\"[notebook] Model not found in S3, downloading from HuggingFace...\")\n", - " from huggingface_hub import snapshot_download\n", - " \n", - " token = os.getenv(\"HUGGINGFACE_HUB_TOKEN\")\n", - " model_path = snapshot_download(\n", - " repo_id=\"Qwen/Qwen2.5-1.5B-Instruct\",\n", - " local_dir=MODEL_DIR,\n", - " token=token,\n", - " resume_download=True,\n", - " local_dir_use_symlinks=False,\n", - " )\n", - " print(f\"✓ Model downloaded to: {model_path}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Training configuration\n", - "MODEL_NAME = \"Qwen/Qwen2.5-1.5B-Instruct\"\n", - "\n", - "# You can also try these alternatives:\n", - "# MODEL_NAME = \"Qwen/Qwen2.5-3B-Instruct\" # Larger, more capable\n", - "# MODEL_NAME = \"Qwen/Qwen2.5-0.5B-Instruct\" # Smaller, faster training\n", - "# MODEL_NAME = \"meta-llama/Llama-3.2-1B-Instruct\" # Alternative architecture\n", - "\n", - "# LoRA configuration\n", - "LORA_R = 16 # Rank - start small, increase if needed\n", - "LORA_ALPHA = 32 # Alpha - typically 2x rank\n", - "LORA_DROPOUT = 0.0 # Dropout - 0.0 is optimized for Unsloth\n", - "\n", - "# Training configuration\n", - "NUM_EPOCHS = 2 # More epochs = better learning, longer training\n", - "LEARNING_RATE = 2e-4 # Standard LoRA learning rate\n", - "MAX_SEQ_LEN = 1024 # Maximum sequence length\n", - "MICRO_BATCH_SIZE = 16 # Batch size per GPU (reduce if OOM)\n", - "GRADIENT_ACCUMULATION = 4 # Effective batch = micro_batch * grad_accum\n", - "\n", - "# QLoRA settings (set to True to enable 4-bit quantization)\n", - "USE_QLORA = True # Set to True if you have limited GPU memory\n", - "\n", - "params = {\n", - " # Model and data path\n", - " 'model_path': MODEL_NAME,\n", - " 'data_path': \"/opt/app-root/src/txt-sql-data/train/train_All_100.jsonl\",\n", - " 'ckpt_output_dir': \"/opt/app-root/src/checkpoints-logs-dir\",\n", - " 'data_output_path': \"/opt/app-root/src/lora-json/_data\",\n", - " # Important for LORA\n", - " 'lr_scheduler': \"cosine\",\n", - " 'warmup_steps': 0,\n", - " 'seed': 42,\n", - " # LoRA configuration\n", - " 'lora_r': LORA_R,\n", - " 'lora_alpha': LORA_ALPHA,\n", - " 'lora_dropout': LORA_DROPOUT,\n", - "\n", - " # Training configuration\n", - " 'num_epochs': NUM_EPOCHS,\n", - " 'learning_rate': LEARNING_RATE,\n", - " 'micro_batch_size': MICRO_BATCH_SIZE,\n", - " 'max_seq_len': MAX_SEQ_LEN,\n", - " 'gradient_accumulation_steps': GRADIENT_ACCUMULATION,\n", - "\n", - " # Dataset format\n", - " 'dataset_type' : \"chat_template\",\n", - " 'field_messages' : \"messages\",\n", - " # Quantization\n", - " 'load_in_4bit': USE_QLORA,\n", - " #GPU configuration\n", - " 'nproc_per_node' : 2,\n", - " 'nnodes' : 2,\n", - " # Logging\n", - " 'logging_steps': 10,\n", - " 'save_steps': 200,\n", - " 'save_total_limit': 3,\n", - "\n", - " # Model Checkpointing\n", - " 'save_final_checkpoint': True,\n", - " 'checkpoint_at_epoch': 2,\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from kubeflow.trainer import TrainerClient\n", - "from kubeflow.trainer.rhai import TrainingHubAlgorithms\n", - "from kubeflow.trainer.rhai import TrainingHubTrainer\n", - "from kubeflow.common.types import KubernetesBackendConfig\n", - "\n", - "backend_cfg = KubernetesBackendConfig(client_configuration=api_client.configuration)\n", - "client = TrainerClient(backend_cfg)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "training_runtime_name = os.getenv(\"TRAINING_RUNTIME\")\n", - "\n", - "if not training_runtime_name:\n", - " raise RuntimeError(\"TRAINING_RUNTIME environment variable is required\")\n", - "\n", - "th_runtime = None\n", - "for runtime in client.list_runtimes():\n", - " if runtime.name == training_runtime_name:\n", - " th_runtime = runtime\n", - " print(\"Found runtime: \" + str(th_runtime))\n", - " break\n", - "\n", - "if th_runtime is None:\n", - " raise RuntimeError(f\"Required runtime '{training_runtime_name}' not found\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from kubeflow.trainer.options.kubernetes import (\n", - " PodTemplateOverrides,\n", - " PodTemplateOverride,\n", - " PodSpecOverride,\n", - " ContainerOverride,\n", - ")\n", - "\n", - "cache_root = \"/opt/app-root/src/.cache/huggingface\"\n", - "triton_cache = \"/opt/app-root/src/.triton\"\n", - "\n", - "job_name = client.train(\n", - " trainer=TrainingHubTrainer(\n", - " algorithm=TrainingHubAlgorithms.LORA_SFT,\n", - " func_args=params,\n", - " env={\n", - " \"HF_HOME\": cache_root,\n", - " \"TRITON_CACHE_DIR\": triton_cache,\n", - " \"XDG_CACHE_HOME\": \"/opt/app-root/src/.cache\",\n", - " \"NCCL_DEBUG\": \"INFO\",\n", - " },\n", - " resources_per_node={\n", - " \"cpu\": 4,\n", - " \"memory\": \"32Gi\",\n", - " \"nvidia.com/gpu\": 1\n", - " },\n", - " ),\n", - " options=[\n", - " PodTemplateOverrides(\n", - " PodTemplateOverride(\n", - " target_jobs=[\"node\"],\n", - " spec=PodSpecOverride(\n", - " volumes=[\n", - " {\"name\": \"work\", \"persistentVolumeClaim\": {\"claimName\": PVC_NAME}},\n", - " ],\n", - " containers=[\n", - " ContainerOverride(\n", - " name=\"node\",\n", - " volume_mounts=[\n", - " {\"name\": \"work\", \"mountPath\": \"/opt/app-root/src\", \"readOnly\": False},\n", - " ],\n", - " )\n", - " ],\n", - " ),\n", - " )\n", - " )\n", - " ],\n", - " runtime=th_runtime,\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Wait for the running status, then wait for completion or failure\n", - "# Using reasonable timeout for LORA training\n", - "client.wait_for_job_status(name=job_name, status={\"Running\"}, timeout=300)\n", - "client.wait_for_job_status(name=job_name, status={\"Complete\", \"Failed\"}, timeout=1800) # 30 minutes for training\n", - "\n", - "# Get job details and logs\n", - "job = client.get_job(name=job_name)\n", - "pod_logs = client.get_job_logs(name=job_name, follow=False)\n", - "\n", - "# Collect all log lines from the generator into a list\n", - "logs = list(pod_logs)\n", - "log_text = \"\\n\".join(str(line) for line in logs)\n", - "\n", - "print(f\"Training job final status: {job.status}\")\n", - "\n", - "# Check 1: Job status must not be \"Failed\" \n", - "if job.status == \"Failed\":\n", - " print(f\"ERROR: Training job '{job_name}' has Failed status\")\n", - " print(\"Last 30 lines of logs:\")\n", - " for line in logs[-30:]:\n", - " print(line)\n", - " raise RuntimeError(f\"Training job '{job_name}' failed\")\n", - "\n", - "# Check 2: Look for the training completion message in logs\n", - "# This is critical because the training script may catch exceptions and exit 0\n", - "if \"[PY] LORA_SFT training complete. Result=\" not in log_text:\n", - " print(f\"ERROR: Training completion message not found in logs\")\n", - " print(\"Last 50 lines of logs:\")\n", - " for line in logs[-50:]:\n", - " print(line)\n", - " raise RuntimeError(f\"Training did not complete successfully - missing completion message\")\n", - "\n", - "print(f\"✓ Training job '{job_name}' completed successfully\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "for c in client.get_job(name=job_name).steps:\n", - " print(f\"Step: {c.name}, Status: {c.status}, Devices: {c.device} x {c.device_count}\\n\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "logs = client.get_job_logs(name=job_name, follow=False)\n", - "\n", - "# Collect all log lines from the generator into a list\n", - "logs = list(pod_logs)\n", - "log_text = \"\\n\".join(str(line) for line in logs)\n", - "print(log_text)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "client.delete_job(job_name)" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": ".venv", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.14.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "nbformat": 4, + "nbformat_minor": 4 } diff --git a/tests/trainer/resources/osft.ipynb b/tests/trainer/resources/osft.ipynb index c1af48d36..2911cff55 100644 --- a/tests/trainer/resources/osft.ipynb +++ b/tests/trainer/resources/osft.ipynb @@ -1,571 +1,570 @@ { - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "%pip install datasets\n", - "\n", - "# pip Install kubeflow SDK from main branch for testing\n", - "%pip install git+https://github.com/opendatahub-io/kubeflow-sdk.git@main" - ] + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%pip install datasets\n", + "\n", + "# kubeflow SDK is installed by test harness via install_kubeflow.py" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "# Standard library imports\n", + "import logging\n", + "import os\n", + "import sys\n", + "import time\n", + "from io import StringIO" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from kubernetes import client as k8s, config as k8s_config\n", + "# Edit to match your specific settings\n", + "api_server = os.getenv(\"OPENSHIFT_API_URL\")\n", + "token = os.getenv(\"NOTEBOOK_USER_TOKEN\")\n", + "if not api_server or not token:\n", + " raise RuntimeError(\"OPENSHIFT_API_URL and NOTEBOOK_USER_TOKEN environment variables are required\")\n", + "PVC_NAME = os.getenv(\"SHARED_PVC_NAME\", \"shared\")\n", + "\n", + "configuration = k8s.Configuration()\n", + "configuration.host = api_server\n", + "# Un-comment if your cluster API server uses a self-signed certificate or an un-trusted CA\n", + "configuration.verify_ssl = False\n", + "configuration.api_key = {\"authorization\": f\"Bearer {token}\"}\n", + "api_client = k8s.ApiClient(configuration)\n", + "\n", + "PVC_MOUNT_PATH = \"/opt/app-root/src\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import gzip\n", + "import shutil\n", + "import socket\n", + "\n", + "import boto3\n", + "from botocore.config import Config as BotoConfig\n", + "from botocore.exceptions import ClientError\n", + "\n", + "# --- Global networking safety net: cap all socket operations ---\n", + "socket.setdefaulttimeout(10) # seconds\n", + "\n", + "# Notebook's PVC mount path (per Notebook CR). Training pods will mount the same PVC at /opt/app-root/src\n", + "PVC_NOTEBOOK_PATH = \"/opt/app-root/src\"\n", + "DATASET_ROOT_NOTEBOOK = PVC_NOTEBOOK_PATH\n", + "TABLE_GPT_DIR = os.path.join(DATASET_ROOT_NOTEBOOK, \"table-gpt-data\", \"train\")\n", + "MODEL_DIR = os.path.join(DATASET_ROOT_NOTEBOOK, \"Qwen\", \"Qwen2.5-1.5B-Instruct\")\n", + "os.makedirs(TABLE_GPT_DIR, exist_ok=True)\n", + "os.makedirs(MODEL_DIR, exist_ok=True)\n", + "\n", + "# Env config for S3/MinIO\n", + "s3_endpoint = os.getenv(\"AWS_DEFAULT_ENDPOINT\", \"\")\n", + "s3_access_key = os.getenv(\"AWS_ACCESS_KEY_ID\", \"\")\n", + "s3_secret_key = os.getenv(\"AWS_SECRET_ACCESS_KEY\", \"\")\n", + "s3_bucket = os.getenv(\"AWS_STORAGE_BUCKET\", \"\")\n", + "s3_prefix = os.getenv(\"AWS_STORAGE_BUCKET_OSFT_DIR\", \"\") # e.g. \"osft-data\"\n", + "\n", + "data_download_successful = False\n", + "\n", + "def stream_download(s3, bucket, key, dst):\n", + " \"\"\"\n", + " Download an object from S3/MinIO using get_object and streaming reads.\n", + " Returns True on success, False on any error.\n", + " \"\"\"\n", + " print(f\"[notebook] STREAM download s3://{bucket}/{key} -> {dst}\")\n", + " t0 = time.time()\n", + "\n", + " try:\n", + " resp = s3.get_object(Bucket=bucket, Key=key)\n", + " except ClientError as e:\n", + " err = e.response.get(\"Error\", {})\n", + " print(f\"[notebook] CLIENT ERROR (get_object) for {key}: {err}\")\n", + " return False\n", + " except Exception as e:\n", + " print(f\"[notebook] OTHER ERROR (get_object) for {key}: {e}\")\n", + " return False\n", + "\n", + " body = resp[\"Body\"]\n", + " try:\n", + " with open(dst, \"wb\") as f:\n", + " while True:\n", + " try:\n", + " chunk = body.read(1024 * 1024) # 1MB per chunk\n", + " except socket.timeout as e:\n", + " print(f\"[notebook] socket.timeout while reading {key}: {e}\")\n", + " return False\n", + " if not chunk:\n", + " break\n", + " f.write(chunk)\n", + " except Exception as e:\n", + " print(f\"[notebook] ERROR writing to {dst} for {key}: {e}\")\n", + " return False\n", + "\n", + " t1 = time.time()\n", + " print(f\"[notebook] DONE stream {key} in {t1 - t0:.2f}s\")\n", + " return True\n", + "\n", + "\n", + "if s3_endpoint and s3_bucket:\n", + " try:\n", + " # Normalize endpoint URL\n", + " endpoint_url = (\n", + " s3_endpoint\n", + " if s3_endpoint.startswith(\"http\")\n", + " else f\"https://{s3_endpoint}\"\n", + " )\n", + " prefix = (s3_prefix or \"\").strip(\"/\")\n", + "\n", + " print(\n", + " f\"[notebook] S3 configured: \"\n", + " f\"endpoint={endpoint_url}, bucket={s3_bucket}, prefix={prefix or ''}\"\n", + " )\n", + "\n", + " # Boto config: single attempt, reasonable connect/read timeouts\n", + " boto_cfg = BotoConfig(\n", + " signature_version=\"s3v4\",\n", + " s3={\"addressing_style\": \"path\"},\n", + " retries={\"max_attempts\": 1, \"mode\": \"standard\"},\n", + " connect_timeout=5,\n", + " read_timeout=10,\n", + " )\n", + "\n", + " # Create S3/MinIO client\n", + " s3 = boto3.client(\n", + " \"s3\",\n", + " endpoint_url=endpoint_url,\n", + " aws_access_key_id=s3_access_key,\n", + " aws_secret_access_key=s3_secret_key,\n", + " config=boto_cfg,\n", + " verify=False,\n", + " )\n", + "\n", + " # List and download all objects under the prefix\n", + " paginator = s3.get_paginator(\"list_objects_v2\")\n", + " pulled_any = False\n", + " file_count = 0\n", + "\n", + " print(f\"[notebook] Starting S3 download from prefix: {prefix}\")\n", + " for page in paginator.paginate(Bucket=s3_bucket, Prefix=prefix or \"\"):\n", + " contents = page.get(\"Contents\", [])\n", + " if not contents:\n", + " print(f\"[notebook] No contents found in this page\")\n", + " continue\n", + " \n", + " print(f\"[notebook] Found {len(contents)} objects in this page\")\n", + "\n", + " for obj in contents:\n", + " key = obj[\"Key\"]\n", + " file_count += 1\n", + "\n", + " # Skip \"directory markers\"\n", + " if key.endswith(\"/\"):\n", + " print(f\"[notebook] Skipping directory marker: {key}\")\n", + " continue\n", + "\n", + " # Determine relative path under prefix for local storage\n", + " rel = key[len(prefix):].lstrip(\"/\") if prefix else key\n", + " print(f\"[notebook] Processing key={key}, rel={rel}\")\n", + " \n", + " # Route to appropriate directory based on content type\n", + " if \"table-gpt\" in rel.lower() or rel.endswith(\".jsonl\"):\n", + " dst = os.path.join(TABLE_GPT_DIR, os.path.basename(rel))\n", + " print(f\"[notebook] Routing to dataset dir: {dst}\")\n", + " elif \"qwen\" in rel.lower() or any(rel.endswith(ext) for ext in [\".bin\", \".json\", \".model\", \".safetensors\", \".txt\"]):\n", + " # Preserve directory structure for model files\n", + " dst = os.path.join(MODEL_DIR, rel.split(\"Qwen2.5-1.5B-Instruct/\")[-1] if \"Qwen2.5-1.5B-Instruct\" in rel else os.path.basename(rel))\n", + " print(f\"[notebook] Routing to model dir: {dst}\")\n", + " else:\n", + " # Default: use the relative path as-is\n", + " dst = os.path.join(DATASET_ROOT_NOTEBOOK, rel)\n", + " print(f\"[notebook] Routing to default dir: {dst}\")\n", + " \n", + " os.makedirs(os.path.dirname(dst), exist_ok=True)\n", + "\n", + " # Download only if missing\n", + " if not os.path.exists(dst):\n", + " ok = stream_download(s3, s3_bucket, key, dst)\n", + " if not ok:\n", + " print(f\"[notebook] Download failed for {key}\")\n", + " continue\n", + " pulled_any = True\n", + " else:\n", + " print(f\"[notebook] Skipping existing file {dst}\")\n", + " pulled_any = True\n", + "\n", + " # If the file is .gz, decompress and remove the .gz\n", + " if dst.endswith(\".gz\") and os.path.exists(dst):\n", + " out_path = os.path.splitext(dst)[0]\n", + " if not os.path.exists(out_path):\n", + " print(f\"[notebook] Decompressing {dst} -> {out_path}\")\n", + " try:\n", + " with gzip.open(dst, \"rb\") as f_in, open(out_path, \"wb\") as f_out:\n", + " shutil.copyfileobj(f_in, f_out)\n", + " except Exception as e:\n", + " print(f\"[notebook] Failed to decompress {dst}: {e}\")\n", + " else:\n", + " try:\n", + " os.remove(dst)\n", + " except Exception:\n", + " pass\n", + "\n", + " if pulled_any:\n", + " print(f\"[notebook] ✓ S3 download successful. Processed {file_count} files\")\n", + " data_download_successful = True\n", + " else:\n", + " print(f\"[notebook] ✗ S3 download found no files to download\")\n", + "\n", + " except Exception as e:\n", + " print(f\"[notebook] ✗ S3 fetch failed: {e}\")\n", + " import traceback\n", + " traceback.print_exc()\n", + " print(\"[notebook] Will attempt HuggingFace fallback...\")\n", + "else:\n", + " print(\"[notebook] S3 not configured (missing endpoint or bucket env vars)\")\n", + "\n", + "# Fallback to HuggingFace if S3 was not configured or failed (requires internet)\n", + "if not data_download_successful:\n", + " print(\"[notebook] Attempting HuggingFace dataset download (requires internet)...\")\n", + " try:\n", + " import json\n", + " import random\n", + " from datasets import load_dataset\n", + "\n", + " # Load the Table-GPT dataset\n", + " print(\"[notebook] Loading Table-GPT dataset from HuggingFace...\")\n", + " dataset = load_dataset(\"LipengCS/Table-GPT\", \"All\")\n", + "\n", + " # Get the training split and create a random subset of 100 samples\n", + " train_data = dataset[\"train\"]\n", + " print(f\"[notebook] Original training set size: {len(train_data)}\")\n", + "\n", + " # Create a random subset of 100 samples\n", + " random.seed(42) # For reproducibility\n", + " subset_indices = random.sample(range(len(train_data)), min(100, len(train_data)))\n", + " subset_data = train_data.select(subset_indices)\n", + "\n", + " print(f\"[notebook] Subset size: {len(subset_data)}\")\n", + "\n", + " # Save the subset to a JSONL file\n", + " output_file = os.path.join(TABLE_GPT_DIR, \"train_All_100.jsonl\")\n", + " with open(output_file, \"w\") as f:\n", + " for example in subset_data:\n", + " f.write(json.dumps(example) + \"\\n\")\n", + "\n", + " print(f\"[notebook] ✓ HuggingFace download successful. Subset saved to {output_file}\")\n", + " data_download_successful = True\n", + "\n", + " except Exception as hf_error:\n", + " print(f\"[notebook] ✗ HuggingFace download failed: {hf_error}\")\n", + " import traceback\n", + " traceback.print_exc()\n", + " raise RuntimeError(\n", + " \"Failed to download dataset from both S3 and HuggingFace. \"\n", + " \"In disconnected environments, ensure S3/MinIO is configured with the required data. \"\n", + " \"In connected environments, check your internet connection and credentials.\"\n", + " ) from hf_error\n", + "\n", + "# Verify dataset file exists\n", + "dataset_file = os.path.join(TABLE_GPT_DIR, \"train_All_100.jsonl\")\n", + "if os.path.exists(dataset_file):\n", + " print(f\"[notebook] ✓ Dataset ready: {dataset_file}\")\n", + "else:\n", + " raise RuntimeError(f\"Dataset file not found: {dataset_file}\")\n", + "\n", + "# Verify model directory has files (model will be downloaded during training if not present)\n", + "if os.path.exists(MODEL_DIR) and os.listdir(MODEL_DIR):\n", + " print(f\"[notebook] ✓ Model files ready in: {MODEL_DIR}\")\n", + " print(f\"[notebook] Model files: {os.listdir(MODEL_DIR)[:5]}...\") # Show first 5 files\n", + "else:\n", + " print(f\"[notebook] Note: Model directory is empty: {MODEL_DIR}\")\n", + " print(\"[notebook] Training will download model from HuggingFace during execution\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Model download - use S3 if available, otherwise HuggingFace\n", + "if os.path.exists(MODEL_DIR) and os.listdir(MODEL_DIR):\n", + " model_path = MODEL_DIR\n", + " print(f\"✓ Using local model from S3: {model_path}\")\n", + "else:\n", + " # Download from HuggingFace\n", + " print(\"[notebook] Model not found in S3, downloading from HuggingFace...\")\n", + " from huggingface_hub import snapshot_download\n", + " \n", + " token = os.getenv(\"HUGGINGFACE_HUB_TOKEN\")\n", + " model_path = snapshot_download(\n", + " repo_id=\"Qwen/Qwen2.5-1.5B-Instruct\",\n", + " local_dir=MODEL_DIR,\n", + " token=token,\n", + " resume_download=True,\n", + " local_dir_use_symlinks=False,\n", + " )\n", + " print(f\"✓ Model downloaded to: {model_path}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Use local model path (downloaded in previous cell if not from S3)\n", + "LOCAL_MODEL_PATH = \"/opt/app-root/src/Qwen/Qwen2.5-1.5B-Instruct\"\n", + "\n", + "params = {\n", + " ###########################################################################\n", + " # 🤖 Model + Data Paths #\n", + " ###########################################################################\n", + " \"model_path\": LOCAL_MODEL_PATH, # Always use local path with actual files\n", + " \"data_path\": \"/opt/app-root/src/table-gpt-data/train/train_All_100.jsonl\",\n", + " \"ckpt_output_dir\": \"/opt/app-root/src/checkpoints-logs-dir\",\n", + " \"data_output_path\": \"/opt/app-root/src/osft-json/_data\",\n", + " ############################################################################\n", + " # 🏋️‍♀️ Training Hyperparameters #\n", + " ############################################################################\n", + " # Important for OSFT\n", + " \"unfreeze_rank_ratio\": 0.25,\n", + " # Standard parameters\n", + " \"effective_batch_size\": 128,\n", + " \"learning_rate\": 5.0e-6,\n", + " \"num_epochs\": 1,\n", + " \"lr_scheduler\": \"cosine\",\n", + " \"warmup_steps\": 0,\n", + " \"seed\": 42,\n", + " ###########################################################################\n", + " # 🏎️ Performance Hyperparameters #\n", + " ###########################################################################\n", + " \"use_liger\": True,\n", + " \"max_tokens_per_gpu\": 32000,\n", + " \"max_seq_len\": 2048,\n", + " ############################################################################\n", + " # 💾 Checkpointing Settings #\n", + " ############################################################################\n", + " # Here we only want to save the very last checkpoint\n", + " \"save_final_checkpoint\": True,\n", + " \"checkpoint_at_epoch\": False,\n", + " ############################################################################\n", + " # 🚀 Distributed Training Configuration #\n", + " ############################################################################\n", + " # Override runtime defaults: 2 nodes with 1 GPU each\n", + " \"nnodes\": 2,\n", + " \"nproc_per_node\": 1,\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from kubeflow.trainer import TrainerClient\n", + "from kubeflow.trainer.rhai import TrainingHubAlgorithms\n", + "from kubeflow.trainer.rhai import TrainingHubTrainer\n", + "from kubeflow_trainer_api import models\n", + "from kubeflow.common.types import KubernetesBackendConfig\n", + "\n", + "backend_cfg = KubernetesBackendConfig(\n", + " client_configuration=api_client.configuration, # <— key part\n", + ")\n", + "\n", + "client = TrainerClient(backend_cfg)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "training_runtime_name = os.getenv(\"TRAINING_RUNTIME\")\n", + "if not training_runtime_name:\n", + " raise RuntimeError(\"TRAINING_RUNTIME environment variable is required\")\n", + "\n", + "th_runtime = None\n", + "for runtime in client.list_runtimes():\n", + " if runtime.name == training_runtime_name:\n", + " th_runtime = runtime\n", + " print(\"Found runtime: \" + str(th_runtime))\n", + " break\n", + "\n", + "if th_runtime is None:\n", + " raise RuntimeError(f\"Required runtime '{training_runtime_name}' not found\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "from kubeflow.trainer.options.kubernetes import (\n", + " PodTemplateOverrides,\n", + " PodTemplateOverride,\n", + " PodSpecOverride,\n", + " ContainerOverride,\n", + ")\n", + "\n", + "cache_root = \"/opt/app-root/src/.cache/huggingface\"\n", + "triton_cache = \"/opt/app-root/src/.triton\"\n", + "\n", + "job_name = client.train(\n", + " trainer=TrainingHubTrainer(\n", + " algorithm=TrainingHubAlgorithms.OSFT,\n", + " func_args=params,\n", + " env={\n", + " \"HF_HOME\": cache_root,\n", + " \"TRITON_CACHE_DIR\": triton_cache,\n", + " \"XDG_CACHE_HOME\": \"/opt/app-root/src/.cache\",\n", + " \"NCCL_DEBUG\": \"INFO\",\n", + " },\n", + " resources_per_node={\n", + " \"cpu\": 4,\n", + " \"memory\": \"32Gi\",\n", + " \"nvidia.com/gpu\": 1\n", + " },\n", + " ),\n", + " options=[\n", + " PodTemplateOverrides(\n", + " PodTemplateOverride(\n", + " target_jobs=[\"node\"],\n", + " spec=PodSpecOverride(\n", + " volumes=[\n", + " {\"name\": \"work\", \"persistentVolumeClaim\": {\"claimName\": PVC_NAME}},\n", + " ],\n", + " containers=[\n", + " ContainerOverride(\n", + " name=\"node\",\n", + " volume_mounts=[\n", + " {\"name\": \"work\", \"mountPath\": \"/opt/app-root/src\", \"readOnly\": False},\n", + " ],\n", + " )\n", + " ],\n", + " ),\n", + " )\n", + " )\n", + " ],\n", + " runtime=th_runtime,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Wait for the running status, then wait for completion or failure\n", + "# Using reasonable timeout for OSFT training\n", + "client.wait_for_job_status(name=job_name, status={\"Running\"}, timeout=300)\n", + "client.wait_for_job_status(name=job_name, status={\"Complete\", \"Failed\"}, timeout=1800) # 30 minutes for training\n", + "\n", + "# Get job details and logs\n", + "job = client.get_job(name=job_name)\n", + "pod_logs = client.get_job_logs(name=job_name, follow=False)\n", + "\n", + "# Collect all log lines from the generator into a list\n", + "logs = list(pod_logs)\n", + "log_text = \"\\n\".join(str(line) for line in logs)\n", + "\n", + "print(f\"Training job final status: {job.status}\")\n", + "\n", + "# Check 1: Job status must not be \"Failed\" \n", + "if job.status == \"Failed\":\n", + " print(f\"ERROR: Training job '{job_name}' has Failed status\")\n", + " print(\"Last 30 lines of logs:\")\n", + " for line in logs[-30:]:\n", + " print(line)\n", + " raise RuntimeError(f\"Training job '{job_name}' failed\")\n", + "\n", + "# Check 2: Look for the training completion message in logs\n", + "# This is critical because the training script may catch exceptions and exit 0\n", + "if \"[PY] OSFT training complete. Result=\" not in log_text:\n", + " print(f\"ERROR: Training completion message not found in logs\")\n", + " print(\"Last 50 lines of logs:\")\n", + " for line in logs[-50:]:\n", + " print(line)\n", + " raise RuntimeError(f\"Training did not complete successfully - missing completion message\")\n", + "\n", + "print(f\"✓ Training job '{job_name}' completed successfully\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "for c in client.get_job(name=job_name).steps:\n", + " print(f\"Step: {c.name}, Status: {c.status}, Devices: {c.device} x {c.device_count}\\n\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "logs = client.get_job_logs(name=job_name, follow=False)\n", + "\n", + "# Collect all log lines from the generator into a list\n", + "logs = list(pod_logs)\n", + "log_text = \"\\n\".join(str(line) for line in logs)\n", + "print(log_text)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "client.delete_job(job_name)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3.12", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.9" + } }, - { - "cell_type": "code", - "execution_count": 3, - "metadata": {}, - "outputs": [], - "source": [ - "# Standard library imports\n", - "import logging\n", - "import os\n", - "import sys\n", - "import time\n", - "from io import StringIO" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from kubernetes import client as k8s, config as k8s_config\n", - "# Edit to match your specific settings\n", - "api_server = os.getenv(\"OPENSHIFT_API_URL\")\n", - "token = os.getenv(\"NOTEBOOK_USER_TOKEN\")\n", - "if not api_server or not token:\n", - " raise RuntimeError(\"OPENSHIFT_API_URL and NOTEBOOK_USER_TOKEN environment variables are required\")\n", - "PVC_NAME = os.getenv(\"SHARED_PVC_NAME\", \"shared\")\n", - "\n", - "configuration = k8s.Configuration()\n", - "configuration.host = api_server\n", - "# Un-comment if your cluster API server uses a self-signed certificate or an un-trusted CA\n", - "configuration.verify_ssl = False\n", - "configuration.api_key = {\"authorization\": f\"Bearer {token}\"}\n", - "api_client = k8s.ApiClient(configuration)\n", - "\n", - "PVC_MOUNT_PATH = \"/opt/app-root/src\"" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", - "import gzip\n", - "import shutil\n", - "import socket\n", - "\n", - "import boto3\n", - "from botocore.config import Config as BotoConfig\n", - "from botocore.exceptions import ClientError\n", - "\n", - "# --- Global networking safety net: cap all socket operations ---\n", - "socket.setdefaulttimeout(10) # seconds\n", - "\n", - "# Notebook's PVC mount path (per Notebook CR). Training pods will mount the same PVC at /opt/app-root/src\n", - "PVC_NOTEBOOK_PATH = \"/opt/app-root/src\"\n", - "DATASET_ROOT_NOTEBOOK = PVC_NOTEBOOK_PATH\n", - "TABLE_GPT_DIR = os.path.join(DATASET_ROOT_NOTEBOOK, \"table-gpt-data\", \"train\")\n", - "MODEL_DIR = os.path.join(DATASET_ROOT_NOTEBOOK, \"Qwen\", \"Qwen2.5-1.5B-Instruct\")\n", - "os.makedirs(TABLE_GPT_DIR, exist_ok=True)\n", - "os.makedirs(MODEL_DIR, exist_ok=True)\n", - "\n", - "# Env config for S3/MinIO\n", - "s3_endpoint = os.getenv(\"AWS_DEFAULT_ENDPOINT\", \"\")\n", - "s3_access_key = os.getenv(\"AWS_ACCESS_KEY_ID\", \"\")\n", - "s3_secret_key = os.getenv(\"AWS_SECRET_ACCESS_KEY\", \"\")\n", - "s3_bucket = os.getenv(\"AWS_STORAGE_BUCKET\", \"\")\n", - "s3_prefix = os.getenv(\"AWS_STORAGE_BUCKET_OSFT_DIR\", \"\") # e.g. \"osft-data\"\n", - "\n", - "data_download_successful = False\n", - "\n", - "def stream_download(s3, bucket, key, dst):\n", - " \"\"\"\n", - " Download an object from S3/MinIO using get_object and streaming reads.\n", - " Returns True on success, False on any error.\n", - " \"\"\"\n", - " print(f\"[notebook] STREAM download s3://{bucket}/{key} -> {dst}\")\n", - " t0 = time.time()\n", - "\n", - " try:\n", - " resp = s3.get_object(Bucket=bucket, Key=key)\n", - " except ClientError as e:\n", - " err = e.response.get(\"Error\", {})\n", - " print(f\"[notebook] CLIENT ERROR (get_object) for {key}: {err}\")\n", - " return False\n", - " except Exception as e:\n", - " print(f\"[notebook] OTHER ERROR (get_object) for {key}: {e}\")\n", - " return False\n", - "\n", - " body = resp[\"Body\"]\n", - " try:\n", - " with open(dst, \"wb\") as f:\n", - " while True:\n", - " try:\n", - " chunk = body.read(1024 * 1024) # 1MB per chunk\n", - " except socket.timeout as e:\n", - " print(f\"[notebook] socket.timeout while reading {key}: {e}\")\n", - " return False\n", - " if not chunk:\n", - " break\n", - " f.write(chunk)\n", - " except Exception as e:\n", - " print(f\"[notebook] ERROR writing to {dst} for {key}: {e}\")\n", - " return False\n", - "\n", - " t1 = time.time()\n", - " print(f\"[notebook] DONE stream {key} in {t1 - t0:.2f}s\")\n", - " return True\n", - "\n", - "\n", - "if s3_endpoint and s3_bucket:\n", - " try:\n", - " # Normalize endpoint URL\n", - " endpoint_url = (\n", - " s3_endpoint\n", - " if s3_endpoint.startswith(\"http\")\n", - " else f\"https://{s3_endpoint}\"\n", - " )\n", - " prefix = (s3_prefix or \"\").strip(\"/\")\n", - "\n", - " print(\n", - " f\"[notebook] S3 configured: \"\n", - " f\"endpoint={endpoint_url}, bucket={s3_bucket}, prefix={prefix or ''}\"\n", - " )\n", - "\n", - " # Boto config: single attempt, reasonable connect/read timeouts\n", - " boto_cfg = BotoConfig(\n", - " signature_version=\"s3v4\",\n", - " s3={\"addressing_style\": \"path\"},\n", - " retries={\"max_attempts\": 1, \"mode\": \"standard\"},\n", - " connect_timeout=5,\n", - " read_timeout=10,\n", - " )\n", - "\n", - " # Create S3/MinIO client\n", - " s3 = boto3.client(\n", - " \"s3\",\n", - " endpoint_url=endpoint_url,\n", - " aws_access_key_id=s3_access_key,\n", - " aws_secret_access_key=s3_secret_key,\n", - " config=boto_cfg,\n", - " verify=False,\n", - " )\n", - "\n", - " # List and download all objects under the prefix\n", - " paginator = s3.get_paginator(\"list_objects_v2\")\n", - " pulled_any = False\n", - " file_count = 0\n", - "\n", - " print(f\"[notebook] Starting S3 download from prefix: {prefix}\")\n", - " for page in paginator.paginate(Bucket=s3_bucket, Prefix=prefix or \"\"):\n", - " contents = page.get(\"Contents\", [])\n", - " if not contents:\n", - " print(f\"[notebook] No contents found in this page\")\n", - " continue\n", - " \n", - " print(f\"[notebook] Found {len(contents)} objects in this page\")\n", - "\n", - " for obj in contents:\n", - " key = obj[\"Key\"]\n", - " file_count += 1\n", - "\n", - " # Skip \"directory markers\"\n", - " if key.endswith(\"/\"):\n", - " print(f\"[notebook] Skipping directory marker: {key}\")\n", - " continue\n", - "\n", - " # Determine relative path under prefix for local storage\n", - " rel = key[len(prefix):].lstrip(\"/\") if prefix else key\n", - " print(f\"[notebook] Processing key={key}, rel={rel}\")\n", - " \n", - " # Route to appropriate directory based on content type\n", - " if \"table-gpt\" in rel.lower() or rel.endswith(\".jsonl\"):\n", - " dst = os.path.join(TABLE_GPT_DIR, os.path.basename(rel))\n", - " print(f\"[notebook] Routing to dataset dir: {dst}\")\n", - " elif \"qwen\" in rel.lower() or any(rel.endswith(ext) for ext in [\".bin\", \".json\", \".model\", \".safetensors\", \".txt\"]):\n", - " # Preserve directory structure for model files\n", - " dst = os.path.join(MODEL_DIR, rel.split(\"Qwen2.5-1.5B-Instruct/\")[-1] if \"Qwen2.5-1.5B-Instruct\" in rel else os.path.basename(rel))\n", - " print(f\"[notebook] Routing to model dir: {dst}\")\n", - " else:\n", - " # Default: use the relative path as-is\n", - " dst = os.path.join(DATASET_ROOT_NOTEBOOK, rel)\n", - " print(f\"[notebook] Routing to default dir: {dst}\")\n", - " \n", - " os.makedirs(os.path.dirname(dst), exist_ok=True)\n", - "\n", - " # Download only if missing\n", - " if not os.path.exists(dst):\n", - " ok = stream_download(s3, s3_bucket, key, dst)\n", - " if not ok:\n", - " print(f\"[notebook] Download failed for {key}\")\n", - " continue\n", - " pulled_any = True\n", - " else:\n", - " print(f\"[notebook] Skipping existing file {dst}\")\n", - " pulled_any = True\n", - "\n", - " # If the file is .gz, decompress and remove the .gz\n", - " if dst.endswith(\".gz\") and os.path.exists(dst):\n", - " out_path = os.path.splitext(dst)[0]\n", - " if not os.path.exists(out_path):\n", - " print(f\"[notebook] Decompressing {dst} -> {out_path}\")\n", - " try:\n", - " with gzip.open(dst, \"rb\") as f_in, open(out_path, \"wb\") as f_out:\n", - " shutil.copyfileobj(f_in, f_out)\n", - " except Exception as e:\n", - " print(f\"[notebook] Failed to decompress {dst}: {e}\")\n", - " else:\n", - " try:\n", - " os.remove(dst)\n", - " except Exception:\n", - " pass\n", - "\n", - " if pulled_any:\n", - " print(f\"[notebook] ✓ S3 download successful. Processed {file_count} files\")\n", - " data_download_successful = True\n", - " else:\n", - " print(f\"[notebook] ✗ S3 download found no files to download\")\n", - "\n", - " except Exception as e:\n", - " print(f\"[notebook] ✗ S3 fetch failed: {e}\")\n", - " import traceback\n", - " traceback.print_exc()\n", - " print(\"[notebook] Will attempt HuggingFace fallback...\")\n", - "else:\n", - " print(\"[notebook] S3 not configured (missing endpoint or bucket env vars)\")\n", - "\n", - "# Fallback to HuggingFace if S3 was not configured or failed (requires internet)\n", - "if not data_download_successful:\n", - " print(\"[notebook] Attempting HuggingFace dataset download (requires internet)...\")\n", - " try:\n", - " import json\n", - " import random\n", - " from datasets import load_dataset\n", - "\n", - " # Load the Table-GPT dataset\n", - " print(\"[notebook] Loading Table-GPT dataset from HuggingFace...\")\n", - " dataset = load_dataset(\"LipengCS/Table-GPT\", \"All\")\n", - "\n", - " # Get the training split and create a random subset of 100 samples\n", - " train_data = dataset[\"train\"]\n", - " print(f\"[notebook] Original training set size: {len(train_data)}\")\n", - "\n", - " # Create a random subset of 100 samples\n", - " random.seed(42) # For reproducibility\n", - " subset_indices = random.sample(range(len(train_data)), min(100, len(train_data)))\n", - " subset_data = train_data.select(subset_indices)\n", - "\n", - " print(f\"[notebook] Subset size: {len(subset_data)}\")\n", - "\n", - " # Save the subset to a JSONL file\n", - " output_file = os.path.join(TABLE_GPT_DIR, \"train_All_100.jsonl\")\n", - " with open(output_file, \"w\") as f:\n", - " for example in subset_data:\n", - " f.write(json.dumps(example) + \"\\n\")\n", - "\n", - " print(f\"[notebook] ✓ HuggingFace download successful. Subset saved to {output_file}\")\n", - " data_download_successful = True\n", - "\n", - " except Exception as hf_error:\n", - " print(f\"[notebook] ✗ HuggingFace download failed: {hf_error}\")\n", - " import traceback\n", - " traceback.print_exc()\n", - " raise RuntimeError(\n", - " \"Failed to download dataset from both S3 and HuggingFace. \"\n", - " \"In disconnected environments, ensure S3/MinIO is configured with the required data. \"\n", - " \"In connected environments, check your internet connection and credentials.\"\n", - " ) from hf_error\n", - "\n", - "# Verify dataset file exists\n", - "dataset_file = os.path.join(TABLE_GPT_DIR, \"train_All_100.jsonl\")\n", - "if os.path.exists(dataset_file):\n", - " print(f\"[notebook] ✓ Dataset ready: {dataset_file}\")\n", - "else:\n", - " raise RuntimeError(f\"Dataset file not found: {dataset_file}\")\n", - "\n", - "# Verify model directory has files (model will be downloaded during training if not present)\n", - "if os.path.exists(MODEL_DIR) and os.listdir(MODEL_DIR):\n", - " print(f\"[notebook] ✓ Model files ready in: {MODEL_DIR}\")\n", - " print(f\"[notebook] Model files: {os.listdir(MODEL_DIR)[:5]}...\") # Show first 5 files\n", - "else:\n", - " print(f\"[notebook] Note: Model directory is empty: {MODEL_DIR}\")\n", - " print(\"[notebook] Training will download model from HuggingFace during execution\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Model download - use S3 if available, otherwise HuggingFace\n", - "if os.path.exists(MODEL_DIR) and os.listdir(MODEL_DIR):\n", - " model_path = MODEL_DIR\n", - " print(f\"✓ Using local model from S3: {model_path}\")\n", - "else:\n", - " # Download from HuggingFace\n", - " print(\"[notebook] Model not found in S3, downloading from HuggingFace...\")\n", - " from huggingface_hub import snapshot_download\n", - " \n", - " token = os.getenv(\"HUGGINGFACE_HUB_TOKEN\")\n", - " model_path = snapshot_download(\n", - " repo_id=\"Qwen/Qwen2.5-1.5B-Instruct\",\n", - " local_dir=MODEL_DIR,\n", - " token=token,\n", - " resume_download=True,\n", - " local_dir_use_symlinks=False,\n", - " )\n", - " print(f\"✓ Model downloaded to: {model_path}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Use local model path (downloaded in previous cell if not from S3)\n", - "LOCAL_MODEL_PATH = \"/opt/app-root/src/Qwen/Qwen2.5-1.5B-Instruct\"\n", - "\n", - "params = {\n", - " ###########################################################################\n", - " # 🤖 Model + Data Paths #\n", - " ###########################################################################\n", - " \"model_path\": LOCAL_MODEL_PATH, # Always use local path with actual files\n", - " \"data_path\": \"/opt/app-root/src/table-gpt-data/train/train_All_100.jsonl\",\n", - " \"ckpt_output_dir\": \"/opt/app-root/src/checkpoints-logs-dir\",\n", - " \"data_output_path\": \"/opt/app-root/src/osft-json/_data\",\n", - " ############################################################################\n", - " # 🏋️‍♀️ Training Hyperparameters #\n", - " ############################################################################\n", - " # Important for OSFT\n", - " \"unfreeze_rank_ratio\": 0.25,\n", - " # Standard parameters\n", - " \"effective_batch_size\": 128,\n", - " \"learning_rate\": 5.0e-6,\n", - " \"num_epochs\": 1,\n", - " \"lr_scheduler\": \"cosine\",\n", - " \"warmup_steps\": 0,\n", - " \"seed\": 42,\n", - " ###########################################################################\n", - " # 🏎️ Performance Hyperparameters #\n", - " ###########################################################################\n", - " \"use_liger\": True,\n", - " \"max_tokens_per_gpu\": 32000,\n", - " \"max_seq_len\": 2048,\n", - " ############################################################################\n", - " # 💾 Checkpointing Settings #\n", - " ############################################################################\n", - " # Here we only want to save the very last checkpoint\n", - " \"save_final_checkpoint\": True,\n", - " \"checkpoint_at_epoch\": False,\n", - " ############################################################################\n", - " # 🚀 Distributed Training Configuration #\n", - " ############################################################################\n", - " # Override runtime defaults: 2 nodes with 1 GPU each\n", - " \"nnodes\": 2,\n", - " \"nproc_per_node\": 1,\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from kubeflow.trainer import TrainerClient\n", - "from kubeflow.trainer.rhai import TrainingHubAlgorithms\n", - "from kubeflow.trainer.rhai import TrainingHubTrainer\n", - "from kubeflow_trainer_api import models\n", - "from kubeflow.common.types import KubernetesBackendConfig\n", - "\n", - "backend_cfg = KubernetesBackendConfig(\n", - " client_configuration=api_client.configuration, # <— key part\n", - ")\n", - "\n", - "client = TrainerClient(backend_cfg)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "training_runtime_name = os.getenv(\"TRAINING_RUNTIME\")\n", - "if not training_runtime_name:\n", - " raise RuntimeError(\"TRAINING_RUNTIME environment variable is required\")\n", - "\n", - "th_runtime = None\n", - "for runtime in client.list_runtimes():\n", - " if runtime.name == training_runtime_name:\n", - " th_runtime = runtime\n", - " print(\"Found runtime: \" + str(th_runtime))\n", - " break\n", - "\n", - "if th_runtime is None:\n", - " raise RuntimeError(f\"Required runtime '{training_runtime_name}' not found\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "\n", - "from kubeflow.trainer.options.kubernetes import (\n", - " PodTemplateOverrides,\n", - " PodTemplateOverride,\n", - " PodSpecOverride,\n", - " ContainerOverride,\n", - ")\n", - "\n", - "cache_root = \"/opt/app-root/src/.cache/huggingface\"\n", - "triton_cache = \"/opt/app-root/src/.triton\"\n", - "\n", - "job_name = client.train(\n", - " trainer=TrainingHubTrainer(\n", - " algorithm=TrainingHubAlgorithms.OSFT,\n", - " func_args=params,\n", - " env={\n", - " \"HF_HOME\": cache_root,\n", - " \"TRITON_CACHE_DIR\": triton_cache,\n", - " \"XDG_CACHE_HOME\": \"/opt/app-root/src/.cache\",\n", - " \"NCCL_DEBUG\": \"INFO\",\n", - " },\n", - " resources_per_node={\n", - " \"cpu\": 4,\n", - " \"memory\": \"32Gi\",\n", - " \"nvidia.com/gpu\": 1\n", - " },\n", - " ),\n", - " options=[\n", - " PodTemplateOverrides(\n", - " PodTemplateOverride(\n", - " target_jobs=[\"node\"],\n", - " spec=PodSpecOverride(\n", - " volumes=[\n", - " {\"name\": \"work\", \"persistentVolumeClaim\": {\"claimName\": PVC_NAME}},\n", - " ],\n", - " containers=[\n", - " ContainerOverride(\n", - " name=\"node\",\n", - " volume_mounts=[\n", - " {\"name\": \"work\", \"mountPath\": \"/opt/app-root/src\", \"readOnly\": False},\n", - " ],\n", - " )\n", - " ],\n", - " ),\n", - " )\n", - " )\n", - " ],\n", - " runtime=th_runtime,\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Wait for the running status, then wait for completion or failure\n", - "# Using reasonable timeout for OSFT training\n", - "client.wait_for_job_status(name=job_name, status={\"Running\"}, timeout=300)\n", - "client.wait_for_job_status(name=job_name, status={\"Complete\", \"Failed\"}, timeout=1800) # 30 minutes for training\n", - "\n", - "# Get job details and logs\n", - "job = client.get_job(name=job_name)\n", - "pod_logs = client.get_job_logs(name=job_name, follow=False)\n", - "\n", - "# Collect all log lines from the generator into a list\n", - "logs = list(pod_logs)\n", - "log_text = \"\\n\".join(str(line) for line in logs)\n", - "\n", - "print(f\"Training job final status: {job.status}\")\n", - "\n", - "# Check 1: Job status must not be \"Failed\" \n", - "if job.status == \"Failed\":\n", - " print(f\"ERROR: Training job '{job_name}' has Failed status\")\n", - " print(\"Last 30 lines of logs:\")\n", - " for line in logs[-30:]:\n", - " print(line)\n", - " raise RuntimeError(f\"Training job '{job_name}' failed\")\n", - "\n", - "# Check 2: Look for the training completion message in logs\n", - "# This is critical because the training script may catch exceptions and exit 0\n", - "if \"[PY] OSFT training complete. Result=\" not in log_text:\n", - " print(f\"ERROR: Training completion message not found in logs\")\n", - " print(\"Last 50 lines of logs:\")\n", - " for line in logs[-50:]:\n", - " print(line)\n", - " raise RuntimeError(f\"Training did not complete successfully - missing completion message\")\n", - "\n", - "print(f\"✓ Training job '{job_name}' completed successfully\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "for c in client.get_job(name=job_name).steps:\n", - " print(f\"Step: {c.name}, Status: {c.status}, Devices: {c.device} x {c.device_count}\\n\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "logs = client.get_job_logs(name=job_name, follow=False)\n", - "\n", - "# Collect all log lines from the generator into a list\n", - "logs = list(pod_logs)\n", - "log_text = \"\\n\".join(str(line) for line in logs)\n", - "print(log_text)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "client.delete_job(job_name)" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3.12", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.9" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "nbformat": 4, + "nbformat_minor": 4 } diff --git a/tests/trainer/resources/sft.ipynb b/tests/trainer/resources/sft.ipynb index e615bda36..5941e9023 100644 --- a/tests/trainer/resources/sft.ipynb +++ b/tests/trainer/resources/sft.ipynb @@ -1,649 +1,648 @@ { - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "id": "dc710e7a-3a8c-4a4e-b18a-45e370772cdb", - "metadata": {}, - "outputs": [], - "source": [ - "%pip install datasets transformers accelerate bitsandbytes huggingface_hub\n", - "\n", - "# pip Install kubeflow SDK from main branch for testing\n", - "%pip install git+https://github.com/opendatahub-io/kubeflow-sdk.git@main" - ] + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "dc710e7a-3a8c-4a4e-b18a-45e370772cdb", + "metadata": {}, + "outputs": [], + "source": [ + "%pip install datasets transformers accelerate bitsandbytes huggingface_hub\n", + "\n", + "# kubeflow SDK is installed by test harness via install_kubeflow.py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "zolfezkdo1", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "from kubernetes import client as k8s, config as k8s_config\n", + "\n", + "# Edit to match your specific settings\n", + "api_server = os.getenv(\"OPENSHIFT_API_URL\")\n", + "token = os.getenv(\"NOTEBOOK_USER_TOKEN\")\n", + "if not api_server or not token:\n", + " raise RuntimeError(\"OPENSHIFT_API_URL and NOTEBOOK_USER_TOKEN environment variables are required\")\n", + "PVC_NAME = os.getenv(\"SHARED_PVC_NAME\", \"shared\")\n", + "\n", + "configuration = k8s.Configuration()\n", + "configuration.host = api_server\n", + "# Un-comment if your cluster API server uses a self-signed certificate or an un-trusted CA\n", + "configuration.verify_ssl = False\n", + "configuration.api_key = {\"authorization\": f\"Bearer {token}\"}\n", + "api_client = k8s.ApiClient(configuration)\n", + "\n", + "PVC_MOUNT_PATH = \"/opt/app-root/src\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f8e19419-5b07-42a4-9735-d520f6aadc33", + "metadata": {}, + "outputs": [], + "source": [ + "# Import Kubeflow Trainer Clients, Configs\n", + "import json\n", + "\n", + "from kubeflow.trainer import TrainerClient\n", + "from kubeflow.trainer.rhai import TrainingHubAlgorithms, TrainingHubTrainer\n", + "from kubeflow.common.types import KubernetesBackendConfig\n", + "\n", + "backend_cfg = KubernetesBackendConfig(\n", + " client_configuration=api_client.configuration,\n", + ")\n", + "\n", + "client = TrainerClient(backend_cfg)\n", + "print(client)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0e0f43c4-8685-409c-814e-84688725d7ef", + "metadata": {}, + "outputs": [], + "source": [ + "# S3/MinIO and dataset download\n", + "import os\n", + "import gzip\n", + "import shutil\n", + "import socket\n", + "import time\n", + "import json\n", + "import random\n", + "\n", + "import boto3\n", + "from botocore.config import Config as BotoConfig\n", + "from botocore.exceptions import ClientError\n", + "\n", + "# --- Global networking safety net: cap all socket operations ---\n", + "socket.setdefaulttimeout(10) # seconds\n", + "\n", + "# Notebook's PVC mount path\n", + "PVC_NOTEBOOK_PATH = \"/opt/app-root/src\"\n", + "DATASET_ROOT_NOTEBOOK = PVC_NOTEBOOK_PATH\n", + "TABLE_GPT_DIR = os.path.join(DATASET_ROOT_NOTEBOOK, \"table-gpt-data\", \"train\")\n", + "MODEL_DIR = os.path.join(DATASET_ROOT_NOTEBOOK, \"Qwen\", \"Qwen2.5-1.5B-Instruct\")\n", + "os.makedirs(TABLE_GPT_DIR, exist_ok=True)\n", + "os.makedirs(MODEL_DIR, exist_ok=True)\n", + "\n", + "# Env config for S3/MinIO\n", + "s3_endpoint = os.getenv(\"AWS_DEFAULT_ENDPOINT\", \"\")\n", + "s3_access_key = os.getenv(\"AWS_ACCESS_KEY_ID\", \"\")\n", + "s3_secret_key = os.getenv(\"AWS_SECRET_ACCESS_KEY\", \"\")\n", + "s3_bucket = os.getenv(\"AWS_STORAGE_BUCKET\", \"\")\n", + "s3_prefix = os.getenv(\"AWS_STORAGE_BUCKET_SFT_DIR\", \"\")\n", + "\n", + "data_download_successful = False\n", + "\n", + "def stream_download(s3, bucket, key, dst):\n", + " \"\"\"Download an object from S3/MinIO using streaming reads.\"\"\"\n", + " print(f\"[notebook] STREAM download s3://{bucket}/{key} -> {dst}\")\n", + " t0 = time.time()\n", + " try:\n", + " resp = s3.get_object(Bucket=bucket, Key=key)\n", + " except ClientError as e:\n", + " print(f\"[notebook] CLIENT ERROR for {key}: {e.response.get('Error', {})}\")\n", + " return False\n", + " except Exception as e:\n", + " print(f\"[notebook] OTHER ERROR for {key}: {e}\")\n", + " return False\n", + "\n", + " body = resp[\"Body\"]\n", + " try:\n", + " with open(dst, \"wb\") as f:\n", + " while True:\n", + " try:\n", + " chunk = body.read(1024 * 1024)\n", + " except socket.timeout as e:\n", + " print(f\"[notebook] socket.timeout while reading {key}: {e}\")\n", + " return False\n", + " if not chunk:\n", + " break\n", + " f.write(chunk)\n", + " except Exception as e:\n", + " print(f\"[notebook] ERROR writing to {dst}: {e}\")\n", + " return False\n", + "\n", + " t1 = time.time()\n", + " print(f\"[notebook] DONE stream {key} in {t1 - t0:.2f}s\")\n", + " return True\n", + "\n", + "# Try S3 download first, fall back to HuggingFace if not configured or fails\n", + "if s3_endpoint and s3_bucket:\n", + " try:\n", + " endpoint_url = s3_endpoint if s3_endpoint.startswith(\"http\") else f\"https://{s3_endpoint}\"\n", + " prefix = (s3_prefix or \"\").strip(\"/\")\n", + " \n", + " print(f\"[notebook] S3 configured: endpoint={endpoint_url}, bucket={s3_bucket}, prefix={prefix or ''}\")\n", + " \n", + " boto_cfg = BotoConfig(\n", + " signature_version=\"s3v4\",\n", + " s3={\"addressing_style\": \"path\"},\n", + " retries={\"max_attempts\": 1, \"mode\": \"standard\"},\n", + " connect_timeout=5,\n", + " read_timeout=10,\n", + " )\n", + " \n", + " # Create S3/MinIO client\n", + " s3 = boto3.client(\n", + " \"s3\",\n", + " endpoint_url=endpoint_url,\n", + " aws_access_key_id=s3_access_key,\n", + " aws_secret_access_key=s3_secret_key,\n", + " config=boto_cfg,\n", + " verify=False,\n", + " )\n", + " \n", + " paginator = s3.get_paginator(\"list_objects_v2\")\n", + " pulled_any = False\n", + " file_count = 0\n", + " \n", + " print(f\"[notebook] Starting S3 download from prefix: {prefix}\")\n", + " for page in paginator.paginate(Bucket=s3_bucket, Prefix=prefix or \"\"):\n", + " contents = page.get(\"Contents\", [])\n", + " if not contents:\n", + " print(f\"[notebook] No contents found in this page\")\n", + " continue\n", + " \n", + " print(f\"[notebook] Found {len(contents)} objects in this page\")\n", + " \n", + " for obj in contents:\n", + " key = obj[\"Key\"]\n", + " file_count += 1\n", + " \n", + " # Skip \"directory markers\"\n", + " if key.endswith(\"/\"):\n", + " print(f\"[notebook] Skipping directory marker: {key}\")\n", + " continue\n", + " \n", + " # Determine relative path under prefix for local storage\n", + " rel = key[len(prefix):].lstrip(\"/\") if prefix else key\n", + " print(f\"[notebook] Processing key={key}, rel={rel}\")\n", + " \n", + " # Route to appropriate directory based on content type\n", + " if \"table-gpt\" in rel.lower() or rel.endswith(\".jsonl\"):\n", + " dst = os.path.join(TABLE_GPT_DIR, os.path.basename(rel))\n", + " print(f\"[notebook] Routing to dataset dir: {dst}\")\n", + " elif \"qwen\" in rel.lower() or any(rel.endswith(ext) for ext in [\".bin\", \".json\", \".model\", \".safetensors\", \".txt\"]):\n", + " dst = os.path.join(MODEL_DIR, rel.split(\"Qwen2.5-1.5B-Instruct/\")[-1] if \"Qwen2.5-1.5B-Instruct\" in rel else os.path.basename(rel))\n", + " print(f\"[notebook] Routing to model dir: {dst}\")\n", + " else:\n", + " dst = os.path.join(DATASET_ROOT_NOTEBOOK, rel)\n", + " print(f\"[notebook] Routing to default dir: {dst}\")\n", + " \n", + " os.makedirs(os.path.dirname(dst), exist_ok=True)\n", + " \n", + " # Download only if missing\n", + " if not os.path.exists(dst):\n", + " ok = stream_download(s3, s3_bucket, key, dst)\n", + " if not ok:\n", + " print(f\"[notebook] Download failed for {key}\")\n", + " continue\n", + " pulled_any = True\n", + " else:\n", + " print(f\"[notebook] Skipping existing file {dst}\")\n", + " pulled_any = True\n", + " \n", + " # If the file is .gz, decompress and remove the .gz\n", + " if dst.endswith(\".gz\") and os.path.exists(dst):\n", + " out_path = os.path.splitext(dst)[0]\n", + " if not os.path.exists(out_path):\n", + " print(f\"[notebook] Decompressing {dst} -> {out_path}\")\n", + " try:\n", + " with gzip.open(dst, \"rb\") as f_in, open(out_path, \"wb\") as f_out:\n", + " shutil.copyfileobj(f_in, f_out)\n", + " except Exception as e:\n", + " print(f\"[notebook] Failed to decompress {dst}: {e}\")\n", + " else:\n", + " try:\n", + " os.remove(dst)\n", + " except Exception:\n", + " pass\n", + " \n", + " if pulled_any:\n", + " print(f\"[notebook] ✓ S3 download successful. Processed {file_count} files\")\n", + " data_download_successful = True\n", + " else:\n", + " print(f\"[notebook] ✗ S3 download found no files to download\")\n", + " \n", + " except Exception as e:\n", + " print(f\"[notebook] ✗ S3 fetch failed: {e}\")\n", + " import traceback\n", + " traceback.print_exc()\n", + " print(\"[notebook] Will attempt HuggingFace fallback...\")\n", + "else:\n", + " print(\"[notebook] S3 not configured (missing endpoint or bucket env vars)\")\n", + "\n", + "# Fallback to HuggingFace if S3 was not configured or failed (requires internet)\n", + "if not data_download_successful:\n", + " print(\"[notebook] Attempting HuggingFace dataset download (requires internet)...\")\n", + " try:\n", + " from datasets import load_dataset\n", + " \n", + " print(\"[notebook] Loading Table-GPT dataset from HuggingFace...\")\n", + " dataset = load_dataset(\"LipengCS/Table-GPT\", \"All\")\n", + " \n", + " train_data = dataset[\"train\"]\n", + " print(f\"[notebook] Original training set size: {len(train_data)}\")\n", + " \n", + " # Create a random subset of 100 samples\n", + " random.seed(42)\n", + " subset_indices = random.sample(range(len(train_data)), min(100, len(train_data)))\n", + " subset_data = train_data.select(subset_indices)\n", + " \n", + " print(f\"[notebook] Subset size: {len(subset_data)}\")\n", + " \n", + " # Save the subset to a JSONL file\n", + " output_file = os.path.join(TABLE_GPT_DIR, \"train_All_100.jsonl\")\n", + " with open(output_file, \"w\") as f:\n", + " for example in subset_data:\n", + " f.write(json.dumps(example) + \"\\n\")\n", + " \n", + " print(f\"[notebook] ✓ HuggingFace download successful. Subset saved to {output_file}\")\n", + " data_download_successful = True\n", + " \n", + " except Exception as hf_error:\n", + " print(f\"[notebook] ✗ HuggingFace download failed: {hf_error}\")\n", + " import traceback\n", + " traceback.print_exc()\n", + " raise RuntimeError(\n", + " \"Failed to download dataset from both S3 and HuggingFace. \"\n", + " \"In disconnected environments, ensure S3/MinIO is configured with the required data. \"\n", + " \"In connected environments, check your internet connection and credentials.\"\n", + " ) from hf_error\n", + "\n", + "# Verify dataset file exists\n", + "dataset_file = os.path.join(TABLE_GPT_DIR, \"train_All_100.jsonl\")\n", + "if os.path.exists(dataset_file):\n", + " print(f\"[notebook] ✓ Dataset ready: {dataset_file}\")\n", + "else:\n", + " raise RuntimeError(f\"Dataset file not found: {dataset_file}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14dce5f4-2bfc-4021-ac94-f19a8b491292", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "from huggingface_hub import snapshot_download\n", + "from huggingface_hub.utils import HfHubHTTPError, RepositoryNotFoundError\n", + "\n", + "\n", + "def download_model_snapshot(\n", + " model_id: str,\n", + " output_dir: str,\n", + " revision: str = \"main\",\n", + " token: str | bool | None = None,\n", + " allow_patterns: list[str] | None = None,\n", + " ignore_patterns: list[str] | None = None,\n", + ") -> str:\n", + " \"\"\"\n", + " Downloads a model snapshot from the Hugging Face Hub using print statements.\n", + "\n", + " Args:\n", + " model_id (str): The repository ID of the model (e.g., \"meta-llama/Meta-Llama-3-8B\").\n", + " output_dir (str): The local directory to save the model files.\n", + " revision (str): The model revision (branch, tag, or commit hash).\n", + " token (str | bool | None): Hugging Face auth token.\n", + " - If None (default): Uses HUGGINGFACE_HUB_TOKEN env var or login.\n", + " - If True: Explicitly uses the stored token.\n", + " - If str: Uses the provided token string.\n", + " allow_patterns (list[str], optional): List of glob patterns to include.\n", + " ignore_patterns (list[str], optional): List of glob patterns to exclude.\n", + "\n", + " Returns:\n", + " str: The path to the downloaded snapshot directory.\n", + "\n", + " Raises:\n", + " RepositoryNotFoundError: If the model is not found.\n", + " HfHubHTTPError: For authentication errors (e.g., 401) or other HTTP issues.\n", + " OSError: If the output directory cannot be created.\n", + " Exception: For other unexpected errors.\n", + " \"\"\"\n", + " print(f\"Attempting to download repository: {model_id}\")\n", + " print(f\"Target directory: {output_dir}\")\n", + " print(f\"Revision: {revision}\")\n", + "\n", + " # Ensure the output directory exists\n", + " try:\n", + " os.makedirs(output_dir, exist_ok=True)\n", + " except OSError as e:\n", + " print(f\"Failed to create output directory {output_dir}: {e}\")\n", + " raise e # Re-raise the exception\n", + "\n", + " # Handle token for gated/private models\n", + " if token is None:\n", + " token = os.getenv(\"HUGGINGFACE_HUB_TOKEN\")\n", + " if token:\n", + " print(\n", + " \"Using HUGGINGFACE_HUB_TOKEN environment variable for authentication.\"\n", + " )\n", + " else:\n", + " token = True # Use locally cached token\n", + " print(\"No explicit token. Will use locally cached token (if logged in).\")\n", + "\n", + " try:\n", + " # This is the core command.\n", + " snapshot_path = snapshot_download(\n", + " repo_id=model_id,\n", + " local_dir=output_dir,\n", + " revision=revision,\n", + " token=token,\n", + " allow_patterns=allow_patterns,\n", + " ignore_patterns=ignore_patterns,\n", + " resume_download=True,\n", + " local_dir_use_symlinks=False,\n", + " )\n", + "\n", + " print(f\"✅ Successfully downloaded model to: {snapshot_path}\")\n", + " return snapshot_path\n", + "\n", + " except RepositoryNotFoundError as e:\n", + " print(f\"❌ Error: Repository '{model_id}' (revision: {revision}) not found.\")\n", + " print(\"Please check the model ID and revision name.\")\n", + " raise e\n", + "\n", + " except HfHubHTTPError as e:\n", + " print(f\"❌ HTTP Error: Failed to access {model_id}.\")\n", + " print(f\"Status Code: {e.response.status_code}\")\n", + " if e.response.status_code == 401:\n", + " print(\"This is an 'Unauthorized' error. The model may be private or gated.\")\n", + " print(\"Ensure you have accepted the license terms on the model's HF page.\")\n", + " print(\n", + " \"And that you are authenticated (use `huggingface-cli login` in your terminal).\"\n", + " )\n", + " raise e\n", + "\n", + " except Exception as e:\n", + " print(f\"❌ An unexpected error occurred: {e}\")\n", + " raise e" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7809dab7-ca08-4bb6-8254-2c3b9aed94f7", + "metadata": {}, + "outputs": [], + "source": [ + "# Model download - use S3 if available, otherwise HuggingFace\n", + "if os.path.exists(MODEL_DIR) and os.listdir(MODEL_DIR):\n", + " model_path = MODEL_DIR\n", + " print(f\"✓ Using local model from S3: {model_path}\")\n", + "else:\n", + " # Download from HuggingFace\n", + " print(\"[notebook] Model not found in S3, downloading from HuggingFace...\")\n", + " from huggingface_hub import snapshot_download\n", + " \n", + " token = os.getenv(\"HUGGINGFACE_HUB_TOKEN\")\n", + " model_path = snapshot_download(\n", + " repo_id=\"Qwen/Qwen2.5-1.5B-Instruct\",\n", + " local_dir=MODEL_DIR,\n", + " token=token,\n", + " resume_download=True,\n", + " local_dir_use_symlinks=False,\n", + " )\n", + " print(f\"✓ Model downloaded to: {model_path}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12126d5b-e928-47d3-80a8-6ad521f0908c", + "metadata": {}, + "outputs": [], + "source": [ + "# Use local model path (downloaded in previous cell if not from S3)\n", + "LOCAL_MODEL_PATH = \"/opt/app-root/src/Qwen/Qwen2.5-1.5B-Instruct\"\n", + "\n", + "training_parameters = {\n", + " ################################################################################\n", + " # 🤖 Model + Data Paths #\n", + " ################################################################################\n", + " \"model_path\": LOCAL_MODEL_PATH, # Always use local path with actual files\n", + " \"data_path\": \"/opt/app-root/src/table-gpt-data/train/train_All_100.jsonl\",\n", + " \"ckpt_output_dir\": \"/opt/app-root/src/checkpoints\",\n", + " \"data_output_dir\": \"/opt/app-root/src/sft-data\",\n", + " ################################################################################\n", + " # 🏋️‍♀️ Training Hyperparameters #\n", + " ################################################################################\n", + " \"effective_batch_size\": 128,\n", + " \"learning_rate\": 5e-6,\n", + " \"num_epochs\": 1,\n", + " \"lr_scheduler\": \"cosine\",\n", + " \"warmup_steps\": 0,\n", + " \"seed\": 42,\n", + " ################################################################################\n", + " # 🏎️ Performance Hyperparameters #\n", + " ################################################################################\n", + " \"max_tokens_per_gpu\": 10000,\n", + " \"max_seq_len\": 8192,\n", + " ################################################################################\n", + " # 💾 Checkpointing Settings #\n", + " ################################################################################\n", + " \"checkpoint_at_epoch\": False,\n", + " ################################################################################\n", + " # 🚀 Distributed Training Configuration #\n", + " ################################################################################\n", + " # Override runtime defaults: 2 nodes with 1 GPU each\n", + " \"nnodes\": 2,\n", + " \"nproc_per_node\": 1,\n", + "}\n", + "\n", + "print(\"⚙️ Training Hyperparameters\")\n", + "print(\"=\" * 50)\n", + "print(json.dumps(training_parameters, indent=4))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8ab7fdad-32ae-4f90-86ec-e41edf2aacf6", + "metadata": {}, + "outputs": [], + "source": [ + "# Find the TrainingHub runtime\n", + "training_runtime_name = os.getenv(\"TRAINING_RUNTIME\")\n", + "if not training_runtime_name:\n", + " raise RuntimeError(\"TRAINING_RUNTIME environment variable is required\")\n", + "\n", + "th_runtime = None\n", + "for runtime in client.list_runtimes():\n", + " if runtime.name == training_runtime_name:\n", + " th_runtime = runtime\n", + " print(\"Found runtime: \" + str(th_runtime))\n", + " break\n", + "\n", + "if th_runtime is None:\n", + " raise RuntimeError(f\"Required runtime '{training_runtime_name}' not found\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ltl994l22wj", + "metadata": {}, + "outputs": [], + "source": [ + "from kubeflow.trainer.options.kubernetes import (\n", + " PodTemplateOverrides,\n", + " PodTemplateOverride,\n", + " PodSpecOverride,\n", + " ContainerOverride,\n", + ")\n", + "\n", + "cache_root = \"/opt/app-root/src/.cache/huggingface\"\n", + "triton_cache = \"/opt/app-root/src/.triton\"\n", + "\n", + "job_name = client.train(\n", + " trainer=TrainingHubTrainer(\n", + " algorithm=TrainingHubAlgorithms.SFT,\n", + " func_args=training_parameters,\n", + " env={ \n", + " \"HF_HOME\": cache_root,\n", + " \"TRITON_CACHE_DIR\": triton_cache,\n", + " \"XDG_CACHE_HOME\": \"/opt/app-root/src/.cache\",\n", + " \"NCCL_DEBUG\": \"INFO\",\n", + " },\n", + " resources_per_node={\n", + " \"cpu\": 4,\n", + " \"memory\": \"32Gi\",\n", + " \"nvidia.com/gpu\": 1\n", + " },\n", + " ),\n", + " options=[\n", + " PodTemplateOverrides(\n", + " PodTemplateOverride(\n", + " target_jobs=[\"node\"],\n", + " spec=PodSpecOverride(\n", + " volumes=[\n", + " {\"name\": \"work\", \"persistentVolumeClaim\": {\"claimName\": PVC_NAME}},\n", + " ],\n", + " containers=[\n", + " ContainerOverride(\n", + " name=\"node\", \n", + " volume_mounts=[\n", + " {\"name\": \"work\", \"mountPath\": \"/opt/app-root/src\", \"readOnly\": False},\n", + " ],\n", + " )\n", + " ],\n", + " ),\n", + " )\n", + " )\n", + " ],\n", + " runtime=th_runtime,\n", + ")\n", + "\n", + "print(f\"Training job created: {job_name}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "98ca62a4-7f3e-48b3-94a4-144f1d92cb10", + "metadata": {}, + "outputs": [], + "source": [ + "# Wait for the running status, then wait for completion or failure\n", + "client.wait_for_job_status(name=job_name, status={\"Running\"}, timeout=300)\n", + "client.wait_for_job_status(name=job_name, status={\"Complete\", \"Failed\"}, timeout=3600) # 1 hour for SFT training\n", + "\n", + "# Get job details and logs\n", + "job = client.get_job(name=job_name)\n", + "pod_logs = client.get_job_logs(job_name, follow=False)\n", + "\n", + "# Flatten all pod logs into a single list of lines\n", + "logs = []\n", + "for log_line in pod_logs:\n", + " logs.extend(str(log_line).splitlines())\n", + "\n", + "log_text = \"\\n\".join(logs)\n", + "\n", + "print(f\"Training job final status: {job.status}\")\n", + "\n", + "# Check 1: Job status must not be \"Failed\" \n", + "if job.status == \"Failed\":\n", + " print(f\"ERROR: Training job '{job_name}' has Failed status\")\n", + " print(\"Last 30 lines of logs:\")\n", + " for line in logs[-30:]:\n", + " print(line)\n", + " raise RuntimeError(f\"Training job '{job_name}' failed\")\n", + "\n", + "# Check 2: Look for the training completion message in logs\n", + "# This is critical because the training script may catch exceptions and exit 0\n", + "if \"[PY] SFT training complete. Result=\" not in log_text:\n", + " print(f\"ERROR: Training completion message not found in logs\")\n", + " print(\"Last 50 lines of logs:\")\n", + " for line in logs[-50:]:\n", + " print(line)\n", + " raise RuntimeError(f\"Training did not complete successfully - missing completion message\")\n", + "\n", + "print(f\"✓ Training job '{job_name}' completed successfully\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "39871e5c-d0ff-4b89-9d31-75e2d71bae21", + "metadata": {}, + "outputs": [], + "source": [ + "for c in client.get_job(name=job_name).steps:\n", + " print(f\"Step: {c.name}, Status: {c.status}, Devices: {c.device} x {c.device_count}\\n\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "sovsqgypat", + "metadata": {}, + "outputs": [], + "source": [ + "for logline in client.get_job_logs(job_name, follow=False):\n", + " print(logline)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9c1edf2f-6cf1-4cba-8b61-58955eed4543", + "metadata": {}, + "outputs": [], + "source": [ + "client.delete_job(job_name)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3.12", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.9" + } }, - { - "cell_type": "code", - "execution_count": null, - "id": "zolfezkdo1", - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", - "from kubernetes import client as k8s, config as k8s_config\n", - "\n", - "# Edit to match your specific settings\n", - "api_server = os.getenv(\"OPENSHIFT_API_URL\")\n", - "token = os.getenv(\"NOTEBOOK_USER_TOKEN\")\n", - "if not api_server or not token:\n", - " raise RuntimeError(\"OPENSHIFT_API_URL and NOTEBOOK_USER_TOKEN environment variables are required\")\n", - "PVC_NAME = os.getenv(\"SHARED_PVC_NAME\", \"shared\")\n", - "\n", - "configuration = k8s.Configuration()\n", - "configuration.host = api_server\n", - "# Un-comment if your cluster API server uses a self-signed certificate or an un-trusted CA\n", - "configuration.verify_ssl = False\n", - "configuration.api_key = {\"authorization\": f\"Bearer {token}\"}\n", - "api_client = k8s.ApiClient(configuration)\n", - "\n", - "PVC_MOUNT_PATH = \"/opt/app-root/src\"" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "f8e19419-5b07-42a4-9735-d520f6aadc33", - "metadata": {}, - "outputs": [], - "source": [ - "# Import Kubeflow Trainer Clients, Configs\n", - "import json\n", - "\n", - "from kubeflow.trainer import TrainerClient\n", - "from kubeflow.trainer.rhai import TrainingHubAlgorithms, TrainingHubTrainer\n", - "from kubeflow.common.types import KubernetesBackendConfig\n", - "\n", - "backend_cfg = KubernetesBackendConfig(\n", - " client_configuration=api_client.configuration,\n", - ")\n", - "\n", - "client = TrainerClient(backend_cfg)\n", - "print(client)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "0e0f43c4-8685-409c-814e-84688725d7ef", - "metadata": {}, - "outputs": [], - "source": [ - "# S3/MinIO and dataset download\n", - "import os\n", - "import gzip\n", - "import shutil\n", - "import socket\n", - "import time\n", - "import json\n", - "import random\n", - "\n", - "import boto3\n", - "from botocore.config import Config as BotoConfig\n", - "from botocore.exceptions import ClientError\n", - "\n", - "# --- Global networking safety net: cap all socket operations ---\n", - "socket.setdefaulttimeout(10) # seconds\n", - "\n", - "# Notebook's PVC mount path\n", - "PVC_NOTEBOOK_PATH = \"/opt/app-root/src\"\n", - "DATASET_ROOT_NOTEBOOK = PVC_NOTEBOOK_PATH\n", - "TABLE_GPT_DIR = os.path.join(DATASET_ROOT_NOTEBOOK, \"table-gpt-data\", \"train\")\n", - "MODEL_DIR = os.path.join(DATASET_ROOT_NOTEBOOK, \"Qwen\", \"Qwen2.5-1.5B-Instruct\")\n", - "os.makedirs(TABLE_GPT_DIR, exist_ok=True)\n", - "os.makedirs(MODEL_DIR, exist_ok=True)\n", - "\n", - "# Env config for S3/MinIO\n", - "s3_endpoint = os.getenv(\"AWS_DEFAULT_ENDPOINT\", \"\")\n", - "s3_access_key = os.getenv(\"AWS_ACCESS_KEY_ID\", \"\")\n", - "s3_secret_key = os.getenv(\"AWS_SECRET_ACCESS_KEY\", \"\")\n", - "s3_bucket = os.getenv(\"AWS_STORAGE_BUCKET\", \"\")\n", - "s3_prefix = os.getenv(\"AWS_STORAGE_BUCKET_SFT_DIR\", \"\")\n", - "\n", - "data_download_successful = False\n", - "\n", - "def stream_download(s3, bucket, key, dst):\n", - " \"\"\"Download an object from S3/MinIO using streaming reads.\"\"\"\n", - " print(f\"[notebook] STREAM download s3://{bucket}/{key} -> {dst}\")\n", - " t0 = time.time()\n", - " try:\n", - " resp = s3.get_object(Bucket=bucket, Key=key)\n", - " except ClientError as e:\n", - " print(f\"[notebook] CLIENT ERROR for {key}: {e.response.get('Error', {})}\")\n", - " return False\n", - " except Exception as e:\n", - " print(f\"[notebook] OTHER ERROR for {key}: {e}\")\n", - " return False\n", - "\n", - " body = resp[\"Body\"]\n", - " try:\n", - " with open(dst, \"wb\") as f:\n", - " while True:\n", - " try:\n", - " chunk = body.read(1024 * 1024)\n", - " except socket.timeout as e:\n", - " print(f\"[notebook] socket.timeout while reading {key}: {e}\")\n", - " return False\n", - " if not chunk:\n", - " break\n", - " f.write(chunk)\n", - " except Exception as e:\n", - " print(f\"[notebook] ERROR writing to {dst}: {e}\")\n", - " return False\n", - "\n", - " t1 = time.time()\n", - " print(f\"[notebook] DONE stream {key} in {t1 - t0:.2f}s\")\n", - " return True\n", - "\n", - "# Try S3 download first, fall back to HuggingFace if not configured or fails\n", - "if s3_endpoint and s3_bucket:\n", - " try:\n", - " endpoint_url = s3_endpoint if s3_endpoint.startswith(\"http\") else f\"https://{s3_endpoint}\"\n", - " prefix = (s3_prefix or \"\").strip(\"/\")\n", - " \n", - " print(f\"[notebook] S3 configured: endpoint={endpoint_url}, bucket={s3_bucket}, prefix={prefix or ''}\")\n", - " \n", - " boto_cfg = BotoConfig(\n", - " signature_version=\"s3v4\",\n", - " s3={\"addressing_style\": \"path\"},\n", - " retries={\"max_attempts\": 1, \"mode\": \"standard\"},\n", - " connect_timeout=5,\n", - " read_timeout=10,\n", - " )\n", - " \n", - " # Create S3/MinIO client\n", - " s3 = boto3.client(\n", - " \"s3\",\n", - " endpoint_url=endpoint_url,\n", - " aws_access_key_id=s3_access_key,\n", - " aws_secret_access_key=s3_secret_key,\n", - " config=boto_cfg,\n", - " verify=False,\n", - " )\n", - " \n", - " paginator = s3.get_paginator(\"list_objects_v2\")\n", - " pulled_any = False\n", - " file_count = 0\n", - " \n", - " print(f\"[notebook] Starting S3 download from prefix: {prefix}\")\n", - " for page in paginator.paginate(Bucket=s3_bucket, Prefix=prefix or \"\"):\n", - " contents = page.get(\"Contents\", [])\n", - " if not contents:\n", - " print(f\"[notebook] No contents found in this page\")\n", - " continue\n", - " \n", - " print(f\"[notebook] Found {len(contents)} objects in this page\")\n", - " \n", - " for obj in contents:\n", - " key = obj[\"Key\"]\n", - " file_count += 1\n", - " \n", - " # Skip \"directory markers\"\n", - " if key.endswith(\"/\"):\n", - " print(f\"[notebook] Skipping directory marker: {key}\")\n", - " continue\n", - " \n", - " # Determine relative path under prefix for local storage\n", - " rel = key[len(prefix):].lstrip(\"/\") if prefix else key\n", - " print(f\"[notebook] Processing key={key}, rel={rel}\")\n", - " \n", - " # Route to appropriate directory based on content type\n", - " if \"table-gpt\" in rel.lower() or rel.endswith(\".jsonl\"):\n", - " dst = os.path.join(TABLE_GPT_DIR, os.path.basename(rel))\n", - " print(f\"[notebook] Routing to dataset dir: {dst}\")\n", - " elif \"qwen\" in rel.lower() or any(rel.endswith(ext) for ext in [\".bin\", \".json\", \".model\", \".safetensors\", \".txt\"]):\n", - " dst = os.path.join(MODEL_DIR, rel.split(\"Qwen2.5-1.5B-Instruct/\")[-1] if \"Qwen2.5-1.5B-Instruct\" in rel else os.path.basename(rel))\n", - " print(f\"[notebook] Routing to model dir: {dst}\")\n", - " else:\n", - " dst = os.path.join(DATASET_ROOT_NOTEBOOK, rel)\n", - " print(f\"[notebook] Routing to default dir: {dst}\")\n", - " \n", - " os.makedirs(os.path.dirname(dst), exist_ok=True)\n", - " \n", - " # Download only if missing\n", - " if not os.path.exists(dst):\n", - " ok = stream_download(s3, s3_bucket, key, dst)\n", - " if not ok:\n", - " print(f\"[notebook] Download failed for {key}\")\n", - " continue\n", - " pulled_any = True\n", - " else:\n", - " print(f\"[notebook] Skipping existing file {dst}\")\n", - " pulled_any = True\n", - " \n", - " # If the file is .gz, decompress and remove the .gz\n", - " if dst.endswith(\".gz\") and os.path.exists(dst):\n", - " out_path = os.path.splitext(dst)[0]\n", - " if not os.path.exists(out_path):\n", - " print(f\"[notebook] Decompressing {dst} -> {out_path}\")\n", - " try:\n", - " with gzip.open(dst, \"rb\") as f_in, open(out_path, \"wb\") as f_out:\n", - " shutil.copyfileobj(f_in, f_out)\n", - " except Exception as e:\n", - " print(f\"[notebook] Failed to decompress {dst}: {e}\")\n", - " else:\n", - " try:\n", - " os.remove(dst)\n", - " except Exception:\n", - " pass\n", - " \n", - " if pulled_any:\n", - " print(f\"[notebook] ✓ S3 download successful. Processed {file_count} files\")\n", - " data_download_successful = True\n", - " else:\n", - " print(f\"[notebook] ✗ S3 download found no files to download\")\n", - " \n", - " except Exception as e:\n", - " print(f\"[notebook] ✗ S3 fetch failed: {e}\")\n", - " import traceback\n", - " traceback.print_exc()\n", - " print(\"[notebook] Will attempt HuggingFace fallback...\")\n", - "else:\n", - " print(\"[notebook] S3 not configured (missing endpoint or bucket env vars)\")\n", - "\n", - "# Fallback to HuggingFace if S3 was not configured or failed (requires internet)\n", - "if not data_download_successful:\n", - " print(\"[notebook] Attempting HuggingFace dataset download (requires internet)...\")\n", - " try:\n", - " from datasets import load_dataset\n", - " \n", - " print(\"[notebook] Loading Table-GPT dataset from HuggingFace...\")\n", - " dataset = load_dataset(\"LipengCS/Table-GPT\", \"All\")\n", - " \n", - " train_data = dataset[\"train\"]\n", - " print(f\"[notebook] Original training set size: {len(train_data)}\")\n", - " \n", - " # Create a random subset of 100 samples\n", - " random.seed(42)\n", - " subset_indices = random.sample(range(len(train_data)), min(100, len(train_data)))\n", - " subset_data = train_data.select(subset_indices)\n", - " \n", - " print(f\"[notebook] Subset size: {len(subset_data)}\")\n", - " \n", - " # Save the subset to a JSONL file\n", - " output_file = os.path.join(TABLE_GPT_DIR, \"train_All_100.jsonl\")\n", - " with open(output_file, \"w\") as f:\n", - " for example in subset_data:\n", - " f.write(json.dumps(example) + \"\\n\")\n", - " \n", - " print(f\"[notebook] ✓ HuggingFace download successful. Subset saved to {output_file}\")\n", - " data_download_successful = True\n", - " \n", - " except Exception as hf_error:\n", - " print(f\"[notebook] ✗ HuggingFace download failed: {hf_error}\")\n", - " import traceback\n", - " traceback.print_exc()\n", - " raise RuntimeError(\n", - " \"Failed to download dataset from both S3 and HuggingFace. \"\n", - " \"In disconnected environments, ensure S3/MinIO is configured with the required data. \"\n", - " \"In connected environments, check your internet connection and credentials.\"\n", - " ) from hf_error\n", - "\n", - "# Verify dataset file exists\n", - "dataset_file = os.path.join(TABLE_GPT_DIR, \"train_All_100.jsonl\")\n", - "if os.path.exists(dataset_file):\n", - " print(f\"[notebook] ✓ Dataset ready: {dataset_file}\")\n", - "else:\n", - " raise RuntimeError(f\"Dataset file not found: {dataset_file}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "14dce5f4-2bfc-4021-ac94-f19a8b491292", - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", - "\n", - "from huggingface_hub import snapshot_download\n", - "from huggingface_hub.utils import HfHubHTTPError, RepositoryNotFoundError\n", - "\n", - "\n", - "def download_model_snapshot(\n", - " model_id: str,\n", - " output_dir: str,\n", - " revision: str = \"main\",\n", - " token: str | bool | None = None,\n", - " allow_patterns: list[str] | None = None,\n", - " ignore_patterns: list[str] | None = None,\n", - ") -> str:\n", - " \"\"\"\n", - " Downloads a model snapshot from the Hugging Face Hub using print statements.\n", - "\n", - " Args:\n", - " model_id (str): The repository ID of the model (e.g., \"meta-llama/Meta-Llama-3-8B\").\n", - " output_dir (str): The local directory to save the model files.\n", - " revision (str): The model revision (branch, tag, or commit hash).\n", - " token (str | bool | None): Hugging Face auth token.\n", - " - If None (default): Uses HUGGINGFACE_HUB_TOKEN env var or login.\n", - " - If True: Explicitly uses the stored token.\n", - " - If str: Uses the provided token string.\n", - " allow_patterns (list[str], optional): List of glob patterns to include.\n", - " ignore_patterns (list[str], optional): List of glob patterns to exclude.\n", - "\n", - " Returns:\n", - " str: The path to the downloaded snapshot directory.\n", - "\n", - " Raises:\n", - " RepositoryNotFoundError: If the model is not found.\n", - " HfHubHTTPError: For authentication errors (e.g., 401) or other HTTP issues.\n", - " OSError: If the output directory cannot be created.\n", - " Exception: For other unexpected errors.\n", - " \"\"\"\n", - " print(f\"Attempting to download repository: {model_id}\")\n", - " print(f\"Target directory: {output_dir}\")\n", - " print(f\"Revision: {revision}\")\n", - "\n", - " # Ensure the output directory exists\n", - " try:\n", - " os.makedirs(output_dir, exist_ok=True)\n", - " except OSError as e:\n", - " print(f\"Failed to create output directory {output_dir}: {e}\")\n", - " raise e # Re-raise the exception\n", - "\n", - " # Handle token for gated/private models\n", - " if token is None:\n", - " token = os.getenv(\"HUGGINGFACE_HUB_TOKEN\")\n", - " if token:\n", - " print(\n", - " \"Using HUGGINGFACE_HUB_TOKEN environment variable for authentication.\"\n", - " )\n", - " else:\n", - " token = True # Use locally cached token\n", - " print(\"No explicit token. Will use locally cached token (if logged in).\")\n", - "\n", - " try:\n", - " # This is the core command.\n", - " snapshot_path = snapshot_download(\n", - " repo_id=model_id,\n", - " local_dir=output_dir,\n", - " revision=revision,\n", - " token=token,\n", - " allow_patterns=allow_patterns,\n", - " ignore_patterns=ignore_patterns,\n", - " resume_download=True,\n", - " local_dir_use_symlinks=False,\n", - " )\n", - "\n", - " print(f\"✅ Successfully downloaded model to: {snapshot_path}\")\n", - " return snapshot_path\n", - "\n", - " except RepositoryNotFoundError as e:\n", - " print(f\"❌ Error: Repository '{model_id}' (revision: {revision}) not found.\")\n", - " print(\"Please check the model ID and revision name.\")\n", - " raise e\n", - "\n", - " except HfHubHTTPError as e:\n", - " print(f\"❌ HTTP Error: Failed to access {model_id}.\")\n", - " print(f\"Status Code: {e.response.status_code}\")\n", - " if e.response.status_code == 401:\n", - " print(\"This is an 'Unauthorized' error. The model may be private or gated.\")\n", - " print(\"Ensure you have accepted the license terms on the model's HF page.\")\n", - " print(\n", - " \"And that you are authenticated (use `huggingface-cli login` in your terminal).\"\n", - " )\n", - " raise e\n", - "\n", - " except Exception as e:\n", - " print(f\"❌ An unexpected error occurred: {e}\")\n", - " raise e" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7809dab7-ca08-4bb6-8254-2c3b9aed94f7", - "metadata": {}, - "outputs": [], - "source": [ - "# Model download - use S3 if available, otherwise HuggingFace\n", - "if os.path.exists(MODEL_DIR) and os.listdir(MODEL_DIR):\n", - " model_path = MODEL_DIR\n", - " print(f\"✓ Using local model from S3: {model_path}\")\n", - "else:\n", - " # Download from HuggingFace\n", - " print(\"[notebook] Model not found in S3, downloading from HuggingFace...\")\n", - " from huggingface_hub import snapshot_download\n", - " \n", - " token = os.getenv(\"HUGGINGFACE_HUB_TOKEN\")\n", - " model_path = snapshot_download(\n", - " repo_id=\"Qwen/Qwen2.5-1.5B-Instruct\",\n", - " local_dir=MODEL_DIR,\n", - " token=token,\n", - " resume_download=True,\n", - " local_dir_use_symlinks=False,\n", - " )\n", - " print(f\"✓ Model downloaded to: {model_path}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "12126d5b-e928-47d3-80a8-6ad521f0908c", - "metadata": {}, - "outputs": [], - "source": [ - "# Use local model path (downloaded in previous cell if not from S3)\n", - "LOCAL_MODEL_PATH = \"/opt/app-root/src/Qwen/Qwen2.5-1.5B-Instruct\"\n", - "\n", - "training_parameters = {\n", - " ################################################################################\n", - " # 🤖 Model + Data Paths #\n", - " ################################################################################\n", - " \"model_path\": LOCAL_MODEL_PATH, # Always use local path with actual files\n", - " \"data_path\": \"/opt/app-root/src/table-gpt-data/train/train_All_100.jsonl\",\n", - " \"ckpt_output_dir\": \"/opt/app-root/src/checkpoints\",\n", - " \"data_output_dir\": \"/opt/app-root/src/sft-data\",\n", - " ################################################################################\n", - " # 🏋️‍♀️ Training Hyperparameters #\n", - " ################################################################################\n", - " \"effective_batch_size\": 128,\n", - " \"learning_rate\": 5e-6,\n", - " \"num_epochs\": 1,\n", - " \"lr_scheduler\": \"cosine\",\n", - " \"warmup_steps\": 0,\n", - " \"seed\": 42,\n", - " ################################################################################\n", - " # 🏎️ Performance Hyperparameters #\n", - " ################################################################################\n", - " \"max_tokens_per_gpu\": 10000,\n", - " \"max_seq_len\": 8192,\n", - " ################################################################################\n", - " # 💾 Checkpointing Settings #\n", - " ################################################################################\n", - " \"checkpoint_at_epoch\": False,\n", - " ################################################################################\n", - " # 🚀 Distributed Training Configuration #\n", - " ################################################################################\n", - " # Override runtime defaults: 2 nodes with 1 GPU each\n", - " \"nnodes\": 2,\n", - " \"nproc_per_node\": 1,\n", - "}\n", - "\n", - "print(\"⚙️ Training Hyperparameters\")\n", - "print(\"=\" * 50)\n", - "print(json.dumps(training_parameters, indent=4))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "8ab7fdad-32ae-4f90-86ec-e41edf2aacf6", - "metadata": {}, - "outputs": [], - "source": [ - "# Find the TrainingHub runtime\n", - "training_runtime_name = os.getenv(\"TRAINING_RUNTIME\")\n", - "if not training_runtime_name:\n", - " raise RuntimeError(\"TRAINING_RUNTIME environment variable is required\")\n", - "\n", - "th_runtime = None\n", - "for runtime in client.list_runtimes():\n", - " if runtime.name == training_runtime_name:\n", - " th_runtime = runtime\n", - " print(\"Found runtime: \" + str(th_runtime))\n", - " break\n", - "\n", - "if th_runtime is None:\n", - " raise RuntimeError(f\"Required runtime '{training_runtime_name}' not found\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "ltl994l22wj", - "metadata": {}, - "outputs": [], - "source": [ - "from kubeflow.trainer.options.kubernetes import (\n", - " PodTemplateOverrides,\n", - " PodTemplateOverride,\n", - " PodSpecOverride,\n", - " ContainerOverride,\n", - ")\n", - "\n", - "cache_root = \"/opt/app-root/src/.cache/huggingface\"\n", - "triton_cache = \"/opt/app-root/src/.triton\"\n", - "\n", - "job_name = client.train(\n", - " trainer=TrainingHubTrainer(\n", - " algorithm=TrainingHubAlgorithms.SFT,\n", - " func_args=training_parameters,\n", - " env={ \n", - " \"HF_HOME\": cache_root,\n", - " \"TRITON_CACHE_DIR\": triton_cache,\n", - " \"XDG_CACHE_HOME\": \"/opt/app-root/src/.cache\",\n", - " \"NCCL_DEBUG\": \"INFO\",\n", - " },\n", - " resources_per_node={\n", - " \"cpu\": 4,\n", - " \"memory\": \"32Gi\",\n", - " \"nvidia.com/gpu\": 1\n", - " },\n", - " ),\n", - " options=[\n", - " PodTemplateOverrides(\n", - " PodTemplateOverride(\n", - " target_jobs=[\"node\"],\n", - " spec=PodSpecOverride(\n", - " volumes=[\n", - " {\"name\": \"work\", \"persistentVolumeClaim\": {\"claimName\": PVC_NAME}},\n", - " ],\n", - " containers=[\n", - " ContainerOverride(\n", - " name=\"node\", \n", - " volume_mounts=[\n", - " {\"name\": \"work\", \"mountPath\": \"/opt/app-root/src\", \"readOnly\": False},\n", - " ],\n", - " )\n", - " ],\n", - " ),\n", - " )\n", - " )\n", - " ],\n", - " runtime=th_runtime,\n", - ")\n", - "\n", - "print(f\"Training job created: {job_name}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "98ca62a4-7f3e-48b3-94a4-144f1d92cb10", - "metadata": {}, - "outputs": [], - "source": [ - "# Wait for the running status, then wait for completion or failure\n", - "client.wait_for_job_status(name=job_name, status={\"Running\"}, timeout=300)\n", - "client.wait_for_job_status(name=job_name, status={\"Complete\", \"Failed\"}, timeout=3600) # 1 hour for SFT training\n", - "\n", - "# Get job details and logs\n", - "job = client.get_job(name=job_name)\n", - "pod_logs = client.get_job_logs(job_name, follow=False)\n", - "\n", - "# Flatten all pod logs into a single list of lines\n", - "logs = []\n", - "for log_line in pod_logs:\n", - " logs.extend(str(log_line).splitlines())\n", - "\n", - "log_text = \"\\n\".join(logs)\n", - "\n", - "print(f\"Training job final status: {job.status}\")\n", - "\n", - "# Check 1: Job status must not be \"Failed\" \n", - "if job.status == \"Failed\":\n", - " print(f\"ERROR: Training job '{job_name}' has Failed status\")\n", - " print(\"Last 30 lines of logs:\")\n", - " for line in logs[-30:]:\n", - " print(line)\n", - " raise RuntimeError(f\"Training job '{job_name}' failed\")\n", - "\n", - "# Check 2: Look for the training completion message in logs\n", - "# This is critical because the training script may catch exceptions and exit 0\n", - "if \"[PY] SFT training complete. Result=\" not in log_text:\n", - " print(f\"ERROR: Training completion message not found in logs\")\n", - " print(\"Last 50 lines of logs:\")\n", - " for line in logs[-50:]:\n", - " print(line)\n", - " raise RuntimeError(f\"Training did not complete successfully - missing completion message\")\n", - "\n", - "print(f\"✓ Training job '{job_name}' completed successfully\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "39871e5c-d0ff-4b89-9d31-75e2d71bae21", - "metadata": {}, - "outputs": [], - "source": [ - "for c in client.get_job(name=job_name).steps:\n", - " print(f\"Step: {c.name}, Status: {c.status}, Devices: {c.device} x {c.device_count}\\n\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "sovsqgypat", - "metadata": {}, - "outputs": [], - "source": [ - "for logline in client.get_job_logs(job_name, follow=False):\n", - " print(logline)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "9c1edf2f-6cf1-4cba-8b61-58955eed4543", - "metadata": {}, - "outputs": [], - "source": [ - "client.delete_job(job_name)" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3.12", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.9" - } - }, - "nbformat": 4, - "nbformat_minor": 5 + "nbformat": 4, + "nbformat_minor": 5 } diff --git a/tests/trainer/resources/torchrun_failure.ipynb b/tests/trainer/resources/torchrun_failure.ipynb index 7b0174014..6c2ce7448 100644 --- a/tests/trainer/resources/torchrun_failure.ipynb +++ b/tests/trainer/resources/torchrun_failure.ipynb @@ -1,241 +1,520 @@ { - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "id": "cell-pip-install", - "metadata": {}, - "outputs": [], - "source": [ - "%pip install datasets transformers accelerate bitsandbytes huggingface_hub\n", - "\n", - "# pip Install kubeflow SDK from main branch for testing\n", - "%pip install git+https://github.com/opendatahub-io/kubeflow-sdk.git@main" - ] + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "cell-pip-install", + "metadata": {}, + "outputs": [], + "source": [ + "%pip install datasets transformers accelerate bitsandbytes huggingface_hub\n", + "\n", + "# kubeflow SDK is installed by test harness via install_kubeflow.py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-k8s-setup", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import urllib3\n", + "from kubernetes import client as k8s\n", + "\n", + "# Suppress InsecureRequestWarning since we use verify_ssl = False\n", + "urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)\n", + "\n", + "api_server = os.getenv(\"OPENSHIFT_API_URL\")\n", + "token = os.getenv(\"NOTEBOOK_USER_TOKEN\")\n", + "if not api_server or not token:\n", + " raise RuntimeError(\"OPENSHIFT_API_URL and NOTEBOOK_USER_TOKEN environment variables are required\")\n", + "\n", + "PVC_NAME = os.getenv(\"SHARED_PVC_NAME\", \"shared\")\n", + "\n", + "configuration = k8s.Configuration()\n", + "configuration.host = api_server\n", + "configuration.verify_ssl = False\n", + "configuration.api_key = {\"authorization\": f\"Bearer {token}\"}\n", + "api_client = k8s.ApiClient(configuration)\n", + "\n", + "PVC_MOUNT_PATH = \"/opt/app-root/src\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-trainer-client", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "\n", + "from kubeflow.trainer import TrainerClient\n", + "from kubeflow.trainer.rhai import TrainingHubAlgorithms, TrainingHubTrainer\n", + "from kubeflow.common.types import KubernetesBackendConfig\n", + "\n", + "backend_cfg = KubernetesBackendConfig(\n", + " client_configuration=api_client.configuration,\n", + ")\n", + "\n", + "client = TrainerClient(backend_cfg)\n", + "print(client)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-data-download", + "metadata": {}, + "outputs": [], + "source": [ + "# S3/MinIO and dataset download\n", + "import os\n", + "import gzip\n", + "import shutil\n", + "import time\n", + "import json\n", + "import random\n", + "\n", + "import boto3\n", + "from botocore.config import Config as BotoConfig\n", + "from botocore.exceptions import ClientError\n", + "\n", + "PVC_NOTEBOOK_PATH = \"/opt/app-root/src\"\n", + "DATASET_ROOT_NOTEBOOK = PVC_NOTEBOOK_PATH\n", + "TABLE_GPT_DIR = os.path.join(DATASET_ROOT_NOTEBOOK, \"table-gpt-data\", \"train\")\n", + "MODEL_DIR = os.path.join(DATASET_ROOT_NOTEBOOK, \"Qwen\", \"Qwen2.5-1.5B-Instruct\")\n", + "os.makedirs(TABLE_GPT_DIR, exist_ok=True)\n", + "os.makedirs(MODEL_DIR, exist_ok=True)\n", + "\n", + "s3_endpoint = os.getenv(\"AWS_DEFAULT_ENDPOINT\", \"\")\n", + "s3_access_key = os.getenv(\"AWS_ACCESS_KEY_ID\", \"\")\n", + "s3_secret_key = os.getenv(\"AWS_SECRET_ACCESS_KEY\", \"\")\n", + "s3_bucket = os.getenv(\"AWS_STORAGE_BUCKET\", \"\")\n", + "s3_prefix = os.getenv(\"AWS_STORAGE_BUCKET_SFT_DIR\", \"\")\n", + "\n", + "data_download_successful = False\n", + "\n", + "def stream_download(s3, bucket, key, dst):\n", + " \"\"\"Download an object from S3/MinIO using streaming reads.\"\"\"\n", + " print(f\"[notebook] STREAM download s3://{bucket}/{key} -> {dst}\")\n", + " t0 = time.time()\n", + " try:\n", + " resp = s3.get_object(Bucket=bucket, Key=key)\n", + " except ClientError as e:\n", + " print(f\"[notebook] CLIENT ERROR for {key}: {e.response.get('Error', {})}\")\n", + " return False\n", + " except Exception as e:\n", + " print(f\"[notebook] OTHER ERROR for {key}: {e}\")\n", + " return False\n", + "\n", + " body = resp[\"Body\"]\n", + " try:\n", + " with open(dst, \"wb\") as f:\n", + " while True:\n", + " try:\n", + " chunk = body.read(1024 * 1024)\n", + " except Exception as e:\n", + " print(f\"[notebook] Error while reading {key}: {e}\")\n", + " return False\n", + " if not chunk:\n", + " break\n", + " f.write(chunk)\n", + " except Exception as e:\n", + " print(f\"[notebook] ERROR writing to {dst}: {e}\")\n", + " return False\n", + "\n", + " t1 = time.time()\n", + " print(f\"[notebook] DONE stream {key} in {t1 - t0:.2f}s\")\n", + " return True\n", + "\n", + "if s3_endpoint and s3_bucket:\n", + " try:\n", + " endpoint_url = s3_endpoint if s3_endpoint.startswith(\"http\") else f\"https://{s3_endpoint}\"\n", + " prefix = (s3_prefix or \"\").strip(\"/\")\n", + "\n", + " print(f\"[notebook] S3 configured: endpoint={endpoint_url}, bucket={s3_bucket}, prefix={prefix or ''}\")\n", + "\n", + " boto_cfg = BotoConfig(\n", + " signature_version=\"s3v4\",\n", + " s3={\"addressing_style\": \"path\"},\n", + " retries={\"max_attempts\": 1, \"mode\": \"standard\"},\n", + " connect_timeout=5,\n", + " read_timeout=10,\n", + " )\n", + "\n", + " s3 = boto3.client(\n", + " \"s3\",\n", + " endpoint_url=endpoint_url,\n", + " aws_access_key_id=s3_access_key,\n", + " aws_secret_access_key=s3_secret_key,\n", + " config=boto_cfg,\n", + " verify=False,\n", + " )\n", + "\n", + " paginator = s3.get_paginator(\"list_objects_v2\")\n", + " pulled_any = False\n", + " file_count = 0\n", + "\n", + " print(f\"[notebook] Starting S3 download from prefix: {prefix}\")\n", + " for page in paginator.paginate(Bucket=s3_bucket, Prefix=prefix or \"\"):\n", + " contents = page.get(\"Contents\", [])\n", + " if not contents:\n", + " print(f\"[notebook] No contents found in this page\")\n", + " continue\n", + "\n", + " print(f\"[notebook] Found {len(contents)} objects in this page\")\n", + "\n", + " for obj in contents:\n", + " key = obj[\"Key\"]\n", + " file_count += 1\n", + "\n", + " if key.endswith(\"/\"):\n", + " print(f\"[notebook] Skipping directory marker: {key}\")\n", + " continue\n", + "\n", + " rel = key[len(prefix):].lstrip(\"/\") if prefix else key\n", + " print(f\"[notebook] Processing key={key}, rel={rel}\")\n", + "\n", + " if \"table-gpt\" in rel.lower() or rel.endswith(\".jsonl\"):\n", + " dst = os.path.join(TABLE_GPT_DIR, os.path.basename(rel))\n", + " elif \"qwen\" in rel.lower() or any(rel.endswith(ext) for ext in [\".bin\", \".json\", \".model\", \".safetensors\", \".txt\"]):\n", + " dst = os.path.join(MODEL_DIR, rel.split(\"Qwen2.5-1.5B-Instruct/\")[-1] if \"Qwen2.5-1.5B-Instruct\" in rel else os.path.basename(rel))\n", + " else:\n", + " dst = os.path.join(DATASET_ROOT_NOTEBOOK, rel)\n", + "\n", + " os.makedirs(os.path.dirname(dst), exist_ok=True)\n", + "\n", + " if not os.path.exists(dst):\n", + " ok = stream_download(s3, s3_bucket, key, dst)\n", + " if not ok:\n", + " print(f\"[notebook] Download failed for {key}\")\n", + " continue\n", + " pulled_any = True\n", + " else:\n", + " print(f\"[notebook] Skipping existing file {dst}\")\n", + " pulled_any = True\n", + "\n", + " if dst.endswith(\".gz\") and os.path.exists(dst):\n", + " out_path = os.path.splitext(dst)[0]\n", + " if not os.path.exists(out_path):\n", + " print(f\"[notebook] Decompressing {dst} -> {out_path}\")\n", + " try:\n", + " with gzip.open(dst, \"rb\") as f_in, open(out_path, \"wb\") as f_out:\n", + " shutil.copyfileobj(f_in, f_out)\n", + " except Exception as e:\n", + " print(f\"[notebook] Failed to decompress {dst}: {e}\")\n", + " else:\n", + " try:\n", + " os.remove(dst)\n", + " except Exception:\n", + " pass\n", + "\n", + " if pulled_any:\n", + " print(f\"[notebook] S3 download successful. Processed {file_count} files\")\n", + " data_download_successful = True\n", + " else:\n", + " print(f\"[notebook] S3 download found no files to download\")\n", + "\n", + " except Exception as e:\n", + " print(f\"[notebook] S3 fetch failed: {e}\")\n", + " import traceback\n", + " traceback.print_exc()\n", + " print(\"[notebook] Will attempt HuggingFace fallback...\")\n", + "else:\n", + " print(\"[notebook] S3 not configured (missing endpoint or bucket env vars)\")\n", + "\n", + "if not data_download_successful:\n", + " print(\"[notebook] Attempting HuggingFace dataset download (requires internet)...\")\n", + " try:\n", + " from datasets import load_dataset\n", + "\n", + " print(\"[notebook] Loading Table-GPT dataset from HuggingFace...\")\n", + " dataset = load_dataset(\"LipengCS/Table-GPT\", \"All\")\n", + "\n", + " train_data = dataset[\"train\"]\n", + " print(f\"[notebook] Original training set size: {len(train_data)}\")\n", + "\n", + " random.seed(42)\n", + " subset_indices = random.sample(range(len(train_data)), min(100, len(train_data)))\n", + " subset_data = train_data.select(subset_indices)\n", + "\n", + " print(f\"[notebook] Subset size: {len(subset_data)}\")\n", + "\n", + " output_file = os.path.join(TABLE_GPT_DIR, \"train_All_100.jsonl\")\n", + " with open(output_file, \"w\") as f:\n", + " for example in subset_data:\n", + " f.write(json.dumps(example) + \"\\n\")\n", + "\n", + " print(f\"[notebook] HuggingFace download successful. Subset saved to {output_file}\")\n", + " data_download_successful = True\n", + "\n", + " except Exception as hf_error:\n", + " print(f\"[notebook] HuggingFace download failed: {hf_error}\")\n", + " import traceback\n", + " traceback.print_exc()\n", + " raise RuntimeError(\n", + " \"Failed to download dataset from both S3 and HuggingFace. \"\n", + " \"In disconnected environments, ensure S3/MinIO is configured with the required data. \"\n", + " \"In connected environments, check your internet connection and credentials.\"\n", + " ) from hf_error\n", + "\n", + "dataset_file = os.path.join(TABLE_GPT_DIR, \"train_All_100.jsonl\")\n", + "if os.path.exists(dataset_file):\n", + " print(f\"[notebook] Dataset ready: {dataset_file}\")\n", + "else:\n", + " raise RuntimeError(f\"Dataset file not found: {dataset_file}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-model-download", + "metadata": {}, + "outputs": [], + "source": [ + "# Model download - use S3 if available, otherwise HuggingFace\n", + "if os.path.exists(MODEL_DIR) and os.listdir(MODEL_DIR):\n", + " model_path = MODEL_DIR\n", + " print(f\"Using local model from S3: {model_path}\")\n", + "else:\n", + " print(\"[notebook] Model not found in S3, downloading from HuggingFace...\")\n", + " from huggingface_hub import snapshot_download\n", + "\n", + " token = os.getenv(\"HUGGINGFACE_HUB_TOKEN\")\n", + " model_path = snapshot_download(\n", + " repo_id=\"Qwen/Qwen2.5-1.5B-Instruct\",\n", + " local_dir=MODEL_DIR,\n", + " token=token,\n", + " resume_download=True,\n", + " local_dir_use_symlinks=False,\n", + " )\n", + " print(f\"Model downloaded to: {model_path}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-find-runtime", + "metadata": {}, + "outputs": [], + "source": [ + "training_runtime_name = os.getenv(\"TRAINING_RUNTIME\")\n", + "if not training_runtime_name:\n", + " raise RuntimeError(\"TRAINING_RUNTIME environment variable is required\")\n", + "\n", + "th_runtime = None\n", + "for runtime in client.list_runtimes():\n", + " if runtime.name == training_runtime_name:\n", + " th_runtime = runtime\n", + " print(\"Found runtime: \" + str(th_runtime))\n", + " break\n", + "\n", + "if th_runtime is None:\n", + " raise RuntimeError(f\"Required runtime '{training_runtime_name}' not found\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-submit-oom-job", + "metadata": {}, + "outputs": [], + "source": [ + "# Submit a training job with valid model + data but an extreme max_batch_len\n", + "# value that will cause an OOM crash during the actual torch forward pass.\n", + "# The model loads, data loads, torchrun starts the training loop, and then\n", + "# the first batch allocation exceeds GPU memory.\n", + "\n", + "from kubeflow.trainer.options.kubernetes import (\n", + " PodTemplateOverrides,\n", + " PodTemplateOverride,\n", + " PodSpecOverride,\n", + " ContainerOverride,\n", + ")\n", + "\n", + "LOCAL_MODEL_PATH = \"/opt/app-root/src/Qwen/Qwen2.5-1.5B-Instruct\"\n", + "\n", + "training_parameters = {\n", + " \"model_path\": LOCAL_MODEL_PATH,\n", + " \"data_path\": \"/opt/app-root/src/table-gpt-data/train/train_All_100.jsonl\",\n", + " \"ckpt_output_dir\": \"/opt/app-root/src/checkpoints\",\n", + " \"data_output_dir\": \"/opt/app-root/src/sft-data\",\n", + " \"effective_batch_size\": 128,\n", + " \"learning_rate\": 5e-6,\n", + " \"num_epochs\": 1,\n", + " \"max_seq_len\": 8192,\n", + " # Extreme value: try to pack 10M tokens into a single GPU batch.\n", + " # This will OOM during the forward pass when torch allocates the\n", + " # attention matrices and activations for this massive batch.\n", + " \"max_batch_len\": 10000000,\n", + "}\n", + "\n", + "cache_root = \"/opt/app-root/src/.cache/huggingface\"\n", + "triton_cache = \"/opt/app-root/src/.triton\"\n", + "\n", + "job_name = client.train(\n", + " trainer=TrainingHubTrainer(\n", + " algorithm=TrainingHubAlgorithms.SFT,\n", + " func_args=training_parameters,\n", + " env={\n", + " \"HF_HOME\": cache_root,\n", + " \"TRITON_CACHE_DIR\": triton_cache,\n", + " \"XDG_CACHE_HOME\": \"/opt/app-root/src/.cache\",\n", + " },\n", + " resources_per_node={\n", + " \"cpu\": 4,\n", + " \"memory\": \"32Gi\",\n", + " \"nvidia.com/gpu\": 1,\n", + " },\n", + " ),\n", + " options=[\n", + " PodTemplateOverrides(\n", + " PodTemplateOverride(\n", + " target_jobs=[\"node\"],\n", + " spec=PodSpecOverride(\n", + " volumes=[\n", + " {\"name\": \"work\", \"persistentVolumeClaim\": {\"claimName\": PVC_NAME}},\n", + " ],\n", + " containers=[\n", + " ContainerOverride(\n", + " name=\"node\",\n", + " volume_mounts=[\n", + " {\"name\": \"work\", \"mountPath\": \"/opt/app-root/src\", \"readOnly\": False},\n", + " ],\n", + " )\n", + " ],\n", + " ),\n", + " )\n", + " )\n", + " ],\n", + " runtime=th_runtime,\n", + ")\n", + "\n", + "print(f\"Training job created: {job_name}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-wait-for-failure", + "metadata": {}, + "outputs": [], + "source": [ + "# Wait for the job to start running, then poll for crash evidence.\n", + "# The training pod will start, load the model, begin the training loop,\n", + "# and OOM during the first forward pass.\n", + "# We check for OOM error strings in logs via get_job_logs() and confirm\n", + "# the crash via get_job()/get_job_events() — either the TrainJob reaches\n", + "# Failed status or crash-related events (BackOff, OOMKilled) appear.\n", + "\n", + "import time\n", + "\n", + "NAMESPACE = os.getenv(\"NOTEBOOK_NAMESPACE\")\n", + "\n", + "# Event reasons that indicate a container crash or failure\n", + "FAILURE_EVENT_REASONS = {\"BackOff\", \"CrashLoopBackOff\", \"Failed\", \"OOMKilled\", \"OOMKilling\", \"Killing\"}\n", + "\n", + "try:\n", + " client.wait_for_job_status(name=job_name, status={\"Running\", \"Failed\"}, timeout=300)\n", + "except Exception as e:\n", + " print(f\"Wait for Running/Failed raised: {e}\")\n", + "\n", + "OOM_ERROR_PATTERNS = [\n", + " \"OutOfMemoryError\",\n", + " \"CUDA out of memory\",\n", + " \"torch.OutOfMemoryError\",\n", + " \"torch.cuda.OutOfMemoryError\",\n", + "]\n", + "\n", + "error_found = False\n", + "failure_confirmed = False\n", + "matched_pattern = None\n", + "deadline = time.time() + 600 # 10 minute timeout\n", + "\n", + "while time.time() < deadline:\n", + " # Check logs for OOM error\n", + " if not error_found:\n", + " try:\n", + " log_lines = list(client.get_job_logs(job_name, follow=False))\n", + " log_text = \"\\n\".join(str(line) for line in log_lines)\n", + " for pattern in OOM_ERROR_PATTERNS:\n", + " if pattern in log_text:\n", + " error_found = True\n", + " matched_pattern = pattern\n", + " print(f\"Found expected error in logs: '{pattern}'\")\n", + " break\n", + " except Exception as e:\n", + " print(f\"Log fetch error (retrying): {e}\")\n", + "\n", + " # Check job failure via get_job() status\n", + " if not failure_confirmed:\n", + " try:\n", + " job = client.get_job(name=job_name)\n", + " if job.status == \"Failed\":\n", + " failure_confirmed = True\n", + " print(f\"Job failure confirmed via get_job(): status={job.status}\")\n", + " except Exception as e:\n", + " print(f\"get_job() error (retrying): {e}\")\n", + "\n", + " # Check events via get_job_events() for crash signals\n", + " if not failure_confirmed:\n", + " try:\n", + " events = client.get_job_events(name=job_name)\n", + " for event in events:\n", + " reason = getattr(event, \"reason\", \"\") or \"\"\n", + " if reason in FAILURE_EVENT_REASONS:\n", + " failure_confirmed = True\n", + " message = getattr(event, \"message\", \"\") or \"\"\n", + " print(f\"Job failure confirmed via get_job_events(): reason={reason} message={message[:100]}\")\n", + " break\n", + " except Exception as e:\n", + " print(f\"get_job_events() error (retrying): {e}\")\n", + "\n", + " if error_found and failure_confirmed:\n", + " break\n", + "\n", + " time.sleep(15)\n", + "\n", + "assert error_found, f\"No OOM/crash error found in logs (searched for: {OOM_ERROR_PATTERNS})\"\n", + "assert failure_confirmed, (\n", + " f\"Job failure not confirmed via get_job()/get_job_events() \"\n", + " f\"— final job status: {client.get_job(name=job_name).status}\"\n", + ")\n", + "\n", + "print(f\"Torchrun failure test PASSED: found '{matched_pattern}' in logs and job failure confirmed via SDK\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-cleanup", + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " client.delete_job(job_name)\n", + " print(f\"Deleted job: {job_name}\")\n", + "except Exception as e:\n", + " print(f\"Failed to delete job: {e}\")\n", + "\n", + "print(\"NOTEBOOK_STATUS: SUCCESS\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } }, - { - "cell_type": "code", - "execution_count": null, - "id": "cell-k8s-setup", - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", - "import urllib3\n", - "from kubernetes import client as k8s\n", - "\n", - "# Suppress InsecureRequestWarning since we use verify_ssl = False\n", - "urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)\n", - "\n", - "api_server = os.getenv(\"OPENSHIFT_API_URL\")\n", - "token = os.getenv(\"NOTEBOOK_USER_TOKEN\")\n", - "if not api_server or not token:\n", - " raise RuntimeError(\"OPENSHIFT_API_URL and NOTEBOOK_USER_TOKEN environment variables are required\")\n", - "\n", - "PVC_NAME = os.getenv(\"SHARED_PVC_NAME\", \"shared\")\n", - "\n", - "configuration = k8s.Configuration()\n", - "configuration.host = api_server\n", - "configuration.verify_ssl = False\n", - "configuration.api_key = {\"authorization\": f\"Bearer {token}\"}\n", - "api_client = k8s.ApiClient(configuration)\n", - "\n", - "PVC_MOUNT_PATH = \"/opt/app-root/src\"" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cell-trainer-client", - "metadata": {}, - "outputs": [], - "source": [ - "import json\n", - "\n", - "from kubeflow.trainer import TrainerClient\n", - "from kubeflow.trainer.rhai import TrainingHubAlgorithms, TrainingHubTrainer\n", - "from kubeflow.common.types import KubernetesBackendConfig\n", - "\n", - "backend_cfg = KubernetesBackendConfig(\n", - " client_configuration=api_client.configuration,\n", - ")\n", - "\n", - "client = TrainerClient(backend_cfg)\n", - "print(client)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cell-data-download", - "metadata": {}, - "outputs": [], - "source": "# S3/MinIO and dataset download\nimport os\nimport gzip\nimport shutil\nimport time\nimport json\nimport random\n\nimport boto3\nfrom botocore.config import Config as BotoConfig\nfrom botocore.exceptions import ClientError\n\nPVC_NOTEBOOK_PATH = \"/opt/app-root/src\"\nDATASET_ROOT_NOTEBOOK = PVC_NOTEBOOK_PATH\nTABLE_GPT_DIR = os.path.join(DATASET_ROOT_NOTEBOOK, \"table-gpt-data\", \"train\")\nMODEL_DIR = os.path.join(DATASET_ROOT_NOTEBOOK, \"Qwen\", \"Qwen2.5-1.5B-Instruct\")\nos.makedirs(TABLE_GPT_DIR, exist_ok=True)\nos.makedirs(MODEL_DIR, exist_ok=True)\n\ns3_endpoint = os.getenv(\"AWS_DEFAULT_ENDPOINT\", \"\")\ns3_access_key = os.getenv(\"AWS_ACCESS_KEY_ID\", \"\")\ns3_secret_key = os.getenv(\"AWS_SECRET_ACCESS_KEY\", \"\")\ns3_bucket = os.getenv(\"AWS_STORAGE_BUCKET\", \"\")\ns3_prefix = os.getenv(\"AWS_STORAGE_BUCKET_SFT_DIR\", \"\")\n\ndata_download_successful = False\n\ndef stream_download(s3, bucket, key, dst):\n \"\"\"Download an object from S3/MinIO using streaming reads.\"\"\"\n print(f\"[notebook] STREAM download s3://{bucket}/{key} -> {dst}\")\n t0 = time.time()\n try:\n resp = s3.get_object(Bucket=bucket, Key=key)\n except ClientError as e:\n print(f\"[notebook] CLIENT ERROR for {key}: {e.response.get('Error', {})}\")\n return False\n except Exception as e:\n print(f\"[notebook] OTHER ERROR for {key}: {e}\")\n return False\n\n body = resp[\"Body\"]\n try:\n with open(dst, \"wb\") as f:\n while True:\n try:\n chunk = body.read(1024 * 1024)\n except Exception as e:\n print(f\"[notebook] Error while reading {key}: {e}\")\n return False\n if not chunk:\n break\n f.write(chunk)\n except Exception as e:\n print(f\"[notebook] ERROR writing to {dst}: {e}\")\n return False\n\n t1 = time.time()\n print(f\"[notebook] DONE stream {key} in {t1 - t0:.2f}s\")\n return True\n\nif s3_endpoint and s3_bucket:\n try:\n endpoint_url = s3_endpoint if s3_endpoint.startswith(\"http\") else f\"https://{s3_endpoint}\"\n prefix = (s3_prefix or \"\").strip(\"/\")\n\n print(f\"[notebook] S3 configured: endpoint={endpoint_url}, bucket={s3_bucket}, prefix={prefix or ''}\")\n\n boto_cfg = BotoConfig(\n signature_version=\"s3v4\",\n s3={\"addressing_style\": \"path\"},\n retries={\"max_attempts\": 1, \"mode\": \"standard\"},\n connect_timeout=5,\n read_timeout=10,\n )\n\n s3 = boto3.client(\n \"s3\",\n endpoint_url=endpoint_url,\n aws_access_key_id=s3_access_key,\n aws_secret_access_key=s3_secret_key,\n config=boto_cfg,\n verify=False,\n )\n\n paginator = s3.get_paginator(\"list_objects_v2\")\n pulled_any = False\n file_count = 0\n\n print(f\"[notebook] Starting S3 download from prefix: {prefix}\")\n for page in paginator.paginate(Bucket=s3_bucket, Prefix=prefix or \"\"):\n contents = page.get(\"Contents\", [])\n if not contents:\n print(f\"[notebook] No contents found in this page\")\n continue\n\n print(f\"[notebook] Found {len(contents)} objects in this page\")\n\n for obj in contents:\n key = obj[\"Key\"]\n file_count += 1\n\n if key.endswith(\"/\"):\n print(f\"[notebook] Skipping directory marker: {key}\")\n continue\n\n rel = key[len(prefix):].lstrip(\"/\") if prefix else key\n print(f\"[notebook] Processing key={key}, rel={rel}\")\n\n if \"table-gpt\" in rel.lower() or rel.endswith(\".jsonl\"):\n dst = os.path.join(TABLE_GPT_DIR, os.path.basename(rel))\n elif \"qwen\" in rel.lower() or any(rel.endswith(ext) for ext in [\".bin\", \".json\", \".model\", \".safetensors\", \".txt\"]):\n dst = os.path.join(MODEL_DIR, rel.split(\"Qwen2.5-1.5B-Instruct/\")[-1] if \"Qwen2.5-1.5B-Instruct\" in rel else os.path.basename(rel))\n else:\n dst = os.path.join(DATASET_ROOT_NOTEBOOK, rel)\n\n os.makedirs(os.path.dirname(dst), exist_ok=True)\n\n if not os.path.exists(dst):\n ok = stream_download(s3, s3_bucket, key, dst)\n if not ok:\n print(f\"[notebook] Download failed for {key}\")\n continue\n pulled_any = True\n else:\n print(f\"[notebook] Skipping existing file {dst}\")\n pulled_any = True\n\n if dst.endswith(\".gz\") and os.path.exists(dst):\n out_path = os.path.splitext(dst)[0]\n if not os.path.exists(out_path):\n print(f\"[notebook] Decompressing {dst} -> {out_path}\")\n try:\n with gzip.open(dst, \"rb\") as f_in, open(out_path, \"wb\") as f_out:\n shutil.copyfileobj(f_in, f_out)\n except Exception as e:\n print(f\"[notebook] Failed to decompress {dst}: {e}\")\n else:\n try:\n os.remove(dst)\n except Exception:\n pass\n\n if pulled_any:\n print(f\"[notebook] S3 download successful. Processed {file_count} files\")\n data_download_successful = True\n else:\n print(f\"[notebook] S3 download found no files to download\")\n\n except Exception as e:\n print(f\"[notebook] S3 fetch failed: {e}\")\n import traceback\n traceback.print_exc()\n print(\"[notebook] Will attempt HuggingFace fallback...\")\nelse:\n print(\"[notebook] S3 not configured (missing endpoint or bucket env vars)\")\n\nif not data_download_successful:\n print(\"[notebook] Attempting HuggingFace dataset download (requires internet)...\")\n try:\n from datasets import load_dataset\n\n print(\"[notebook] Loading Table-GPT dataset from HuggingFace...\")\n dataset = load_dataset(\"LipengCS/Table-GPT\", \"All\")\n\n train_data = dataset[\"train\"]\n print(f\"[notebook] Original training set size: {len(train_data)}\")\n\n random.seed(42)\n subset_indices = random.sample(range(len(train_data)), min(100, len(train_data)))\n subset_data = train_data.select(subset_indices)\n\n print(f\"[notebook] Subset size: {len(subset_data)}\")\n\n output_file = os.path.join(TABLE_GPT_DIR, \"train_All_100.jsonl\")\n with open(output_file, \"w\") as f:\n for example in subset_data:\n f.write(json.dumps(example) + \"\\n\")\n\n print(f\"[notebook] HuggingFace download successful. Subset saved to {output_file}\")\n data_download_successful = True\n\n except Exception as hf_error:\n print(f\"[notebook] HuggingFace download failed: {hf_error}\")\n import traceback\n traceback.print_exc()\n raise RuntimeError(\n \"Failed to download dataset from both S3 and HuggingFace. \"\n \"In disconnected environments, ensure S3/MinIO is configured with the required data. \"\n \"In connected environments, check your internet connection and credentials.\"\n ) from hf_error\n\ndataset_file = os.path.join(TABLE_GPT_DIR, \"train_All_100.jsonl\")\nif os.path.exists(dataset_file):\n print(f\"[notebook] Dataset ready: {dataset_file}\")\nelse:\n raise RuntimeError(f\"Dataset file not found: {dataset_file}\")" - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cell-model-download", - "metadata": {}, - "outputs": [], - "source": [ - "# Model download - use S3 if available, otherwise HuggingFace\n", - "if os.path.exists(MODEL_DIR) and os.listdir(MODEL_DIR):\n", - " model_path = MODEL_DIR\n", - " print(f\"Using local model from S3: {model_path}\")\n", - "else:\n", - " print(\"[notebook] Model not found in S3, downloading from HuggingFace...\")\n", - " from huggingface_hub import snapshot_download\n", - "\n", - " token = os.getenv(\"HUGGINGFACE_HUB_TOKEN\")\n", - " model_path = snapshot_download(\n", - " repo_id=\"Qwen/Qwen2.5-1.5B-Instruct\",\n", - " local_dir=MODEL_DIR,\n", - " token=token,\n", - " resume_download=True,\n", - " local_dir_use_symlinks=False,\n", - " )\n", - " print(f\"Model downloaded to: {model_path}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cell-find-runtime", - "metadata": {}, - "outputs": [], - "source": [ - "training_runtime_name = os.getenv(\"TRAINING_RUNTIME\")\n", - "if not training_runtime_name:\n", - " raise RuntimeError(\"TRAINING_RUNTIME environment variable is required\")\n", - "\n", - "th_runtime = None\n", - "for runtime in client.list_runtimes():\n", - " if runtime.name == training_runtime_name:\n", - " th_runtime = runtime\n", - " print(\"Found runtime: \" + str(th_runtime))\n", - " break\n", - "\n", - "if th_runtime is None:\n", - " raise RuntimeError(f\"Required runtime '{training_runtime_name}' not found\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cell-submit-oom-job", - "metadata": {}, - "outputs": [], - "source": [ - "# Submit a training job with valid model + data but an extreme max_batch_len\n", - "# value that will cause an OOM crash during the actual torch forward pass.\n", - "# The model loads, data loads, torchrun starts the training loop, and then\n", - "# the first batch allocation exceeds GPU memory.\n", - "\n", - "from kubeflow.trainer.options.kubernetes import (\n", - " PodTemplateOverrides,\n", - " PodTemplateOverride,\n", - " PodSpecOverride,\n", - " ContainerOverride,\n", - ")\n", - "\n", - "LOCAL_MODEL_PATH = \"/opt/app-root/src/Qwen/Qwen2.5-1.5B-Instruct\"\n", - "\n", - "training_parameters = {\n", - " \"model_path\": LOCAL_MODEL_PATH,\n", - " \"data_path\": \"/opt/app-root/src/table-gpt-data/train/train_All_100.jsonl\",\n", - " \"ckpt_output_dir\": \"/opt/app-root/src/checkpoints\",\n", - " \"data_output_dir\": \"/opt/app-root/src/sft-data\",\n", - " \"effective_batch_size\": 128,\n", - " \"learning_rate\": 5e-6,\n", - " \"num_epochs\": 1,\n", - " \"max_seq_len\": 8192,\n", - " # Extreme value: try to pack 10M tokens into a single GPU batch.\n", - " # This will OOM during the forward pass when torch allocates the\n", - " # attention matrices and activations for this massive batch.\n", - " \"max_batch_len\": 10000000,\n", - "}\n", - "\n", - "cache_root = \"/opt/app-root/src/.cache/huggingface\"\n", - "triton_cache = \"/opt/app-root/src/.triton\"\n", - "\n", - "job_name = client.train(\n", - " trainer=TrainingHubTrainer(\n", - " algorithm=TrainingHubAlgorithms.SFT,\n", - " func_args=training_parameters,\n", - " env={\n", - " \"HF_HOME\": cache_root,\n", - " \"TRITON_CACHE_DIR\": triton_cache,\n", - " \"XDG_CACHE_HOME\": \"/opt/app-root/src/.cache\",\n", - " },\n", - " resources_per_node={\n", - " \"cpu\": 4,\n", - " \"memory\": \"32Gi\",\n", - " \"nvidia.com/gpu\": 1,\n", - " },\n", - " ),\n", - " options=[\n", - " PodTemplateOverrides(\n", - " PodTemplateOverride(\n", - " target_jobs=[\"node\"],\n", - " spec=PodSpecOverride(\n", - " volumes=[\n", - " {\"name\": \"work\", \"persistentVolumeClaim\": {\"claimName\": PVC_NAME}},\n", - " ],\n", - " containers=[\n", - " ContainerOverride(\n", - " name=\"node\",\n", - " volume_mounts=[\n", - " {\"name\": \"work\", \"mountPath\": \"/opt/app-root/src\", \"readOnly\": False},\n", - " ],\n", - " )\n", - " ],\n", - " ),\n", - " )\n", - " )\n", - " ],\n", - " runtime=th_runtime,\n", - ")\n", - "\n", - "print(f\"Training job created: {job_name}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cell-wait-for-failure", - "metadata": {}, - "outputs": [], - "source": "# Wait for the job to start running, then poll for crash evidence.\n# The training pod will start, load the model, begin the training loop,\n# and OOM during the first forward pass.\n# We check for OOM error strings in logs via get_job_logs() and confirm\n# the crash via get_job()/get_job_events() — either the TrainJob reaches\n# Failed status or crash-related events (BackOff, OOMKilled) appear.\n\nimport time\n\nNAMESPACE = os.getenv(\"NOTEBOOK_NAMESPACE\")\n\n# Event reasons that indicate a container crash or failure\nFAILURE_EVENT_REASONS = {\"BackOff\", \"CrashLoopBackOff\", \"Failed\", \"OOMKilled\", \"OOMKilling\", \"Killing\"}\n\ntry:\n client.wait_for_job_status(name=job_name, status={\"Running\", \"Failed\"}, timeout=300)\nexcept Exception as e:\n print(f\"Wait for Running/Failed raised: {e}\")\n\nOOM_ERROR_PATTERNS = [\n \"OutOfMemoryError\",\n \"CUDA out of memory\",\n \"torch.OutOfMemoryError\",\n \"torch.cuda.OutOfMemoryError\",\n]\n\nerror_found = False\nfailure_confirmed = False\nmatched_pattern = None\ndeadline = time.time() + 600 # 10 minute timeout\n\nwhile time.time() < deadline:\n # Check logs for OOM error\n if not error_found:\n try:\n log_lines = list(client.get_job_logs(job_name, follow=False))\n log_text = \"\\n\".join(str(line) for line in log_lines)\n for pattern in OOM_ERROR_PATTERNS:\n if pattern in log_text:\n error_found = True\n matched_pattern = pattern\n print(f\"Found expected error in logs: '{pattern}'\")\n break\n except Exception as e:\n print(f\"Log fetch error (retrying): {e}\")\n\n # Check job failure via get_job() status\n if not failure_confirmed:\n try:\n job = client.get_job(name=job_name)\n if job.status == \"Failed\":\n failure_confirmed = True\n print(f\"Job failure confirmed via get_job(): status={job.status}\")\n except Exception as e:\n print(f\"get_job() error (retrying): {e}\")\n\n # Check events via get_job_events() for crash signals\n if not failure_confirmed:\n try:\n events = client.get_job_events(name=job_name)\n for event in events:\n reason = getattr(event, \"reason\", \"\") or \"\"\n if reason in FAILURE_EVENT_REASONS:\n failure_confirmed = True\n message = getattr(event, \"message\", \"\") or \"\"\n print(f\"Job failure confirmed via get_job_events(): reason={reason} message={message[:100]}\")\n break\n except Exception as e:\n print(f\"get_job_events() error (retrying): {e}\")\n\n if error_found and failure_confirmed:\n break\n\n time.sleep(15)\n\nassert error_found, f\"No OOM/crash error found in logs (searched for: {OOM_ERROR_PATTERNS})\"\nassert failure_confirmed, (\n f\"Job failure not confirmed via get_job()/get_job_events() \"\n f\"— final job status: {client.get_job(name=job_name).status}\"\n)\n\nprint(f\"Torchrun failure test PASSED: found '{matched_pattern}' in logs and job failure confirmed via SDK\")" - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cell-cleanup", - "metadata": {}, - "outputs": [], - "source": [ - "try:\n", - " client.delete_job(job_name)\n", - " print(f\"Deleted job: {job_name}\")\n", - "except Exception as e:\n", - " print(f\"Failed to delete job: {e}\")\n", - "\n", - "print(\"NOTEBOOK_STATUS: SUCCESS\")" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "name": "python", - "version": "3.12.0" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} \ No newline at end of file + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tests/trainer/sdk_tests/failure_traininghub_tests.go b/tests/trainer/sdk_tests/failure_traininghub_tests.go index 47565b002..11608dbaf 100644 --- a/tests/trainer/sdk_tests/failure_traininghub_tests.go +++ b/tests/trainer/sdk_tests/failure_traininghub_tests.go @@ -50,16 +50,21 @@ func RunTrainingFailureScenariosTest(t *testing.T) { trainerutils.EnsureNotebookServiceAccount(t, test, namespace.Name) // RBACs setup - userName := common.GetNotebookUserName(test) - userToken := common.GenerateNotebookUserToken(test) + userToken := common.GetNotebookUserTokenFromEnv(test) + userName := common.GetNotebookUserNameFromEnv(test, userToken) support.CreateUserRoleBindingWithClusterRole(test, userName, namespace.Name, "admin") trainerutils.CreateUserClusterRoleBindingForTrainerRuntimes(test, userName) - // Create ConfigMap with notebook + // Create ConfigMap with notebook and kubeflow install script localPath := failureNotebookPath nb, err := os.ReadFile(localPath) test.Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("failed to read notebook: %s", localPath)) - cm := support.CreateConfigMap(test, namespace.Name, map[string][]byte{failureNotebookName: nb}) + installScript, err := os.ReadFile(installScriptPath) + test.Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("failed to read install script: %s", installScriptPath)) + cm := support.CreateConfigMap(test, namespace.Name, map[string][]byte{ + failureNotebookName: nb, + installKubeflowScript: installScript, + }) // Create RWX PVC required by the notebook pod template storageClass, err := support.GetRWXStorageClass(test) @@ -72,16 +77,21 @@ func RunTrainingFailureScenariosTest(t *testing.T) { support.StorageClassName(storageClass.Name), ) + sdkInstallExports := buildKubeflowInstallExports() shellCmd := fmt.Sprintf( "set -e; "+ "export OPENSHIFT_API_URL='%s'; export NOTEBOOK_USER_TOKEN='%s'; "+ "export NOTEBOOK_NAMESPACE='%s'; "+ "export TRAINING_RUNTIME='%s'; "+ + "%s"+ "python -m pip install --quiet --no-cache-dir --break-system-packages ipykernel papermill && "+ + "python /opt/app-root/notebooks/%s && "+ "if python -m papermill -k python3 /opt/app-root/notebooks/%s /opt/app-root/src/out.ipynb --log-output; "+ "then echo 'NOTEBOOK_STATUS: SUCCESS'; else echo 'NOTEBOOK_STATUS: FAILURE'; fi; sleep infinity", support.GetOpenShiftApiUrl(test), userToken, namespace.Name, trainerutils.DefaultTrainingHubRuntime, + sdkInstallExports, + installKubeflowScript, failureNotebookName, ) command := []string{"/bin/sh", "-c", shellCmd} @@ -117,16 +127,21 @@ func RunTorchrunTrainingFailureTest(t *testing.T) { trainerutils.EnsureNotebookServiceAccount(t, test, namespace.Name) // RBACs setup - userName := common.GetNotebookUserName(test) - userToken := common.GenerateNotebookUserToken(test) + userToken := common.GetNotebookUserTokenFromEnv(test) + userName := common.GetNotebookUserNameFromEnv(test, userToken) support.CreateUserRoleBindingWithClusterRole(test, userName, namespace.Name, "admin") trainerutils.CreateUserClusterRoleBindingForTrainerRuntimes(test, userName) - // Create ConfigMap with notebook + // Create ConfigMap with notebook and kubeflow install script localPath := torchrunFailureNotebookPath nb, err := os.ReadFile(localPath) test.Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("failed to read notebook: %s", localPath)) - cm := support.CreateConfigMap(test, namespace.Name, map[string][]byte{torchrunFailureNotebookName: nb}) + installScript, err := os.ReadFile(installScriptPath) + test.Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("failed to read install script: %s", installScriptPath)) + cm := support.CreateConfigMap(test, namespace.Name, map[string][]byte{ + torchrunFailureNotebookName: nb, + installKubeflowScript: installScript, + }) // S3 configuration for model and dataset download endpoint, endpointOK := support.GetStorageBucketDefaultEndpoint() @@ -152,6 +167,7 @@ func RunTorchrunTrainingFailureTest(t *testing.T) { support.StorageClassName(storageClass.Name), ) + sdkInstallExports := buildKubeflowInstallExports() shellCmd := fmt.Sprintf( "set -e; "+ "export OPENSHIFT_API_URL='%s'; export NOTEBOOK_USER_TOKEN='%s'; "+ @@ -162,12 +178,16 @@ func RunTorchrunTrainingFailureTest(t *testing.T) { "export AWS_STORAGE_BUCKET='%s'; "+ "export AWS_STORAGE_BUCKET_SFT_DIR='%s'; "+ "export TRAINING_RUNTIME='%s'; "+ + "%s"+ "python -m pip install --quiet --no-cache-dir --break-system-packages ipykernel papermill boto3==1.34.162 && "+ + "python /opt/app-root/notebooks/%s && "+ "if python -m papermill -k python3 /opt/app-root/notebooks/%s /opt/app-root/src/out.ipynb --log-output; "+ "then echo 'NOTEBOOK_STATUS: SUCCESS'; else echo 'NOTEBOOK_STATUS: FAILURE'; fi; sleep infinity", support.GetOpenShiftApiUrl(test), userToken, namespace.Name, rwxPvc.Name, endpoint, accessKey, secretKey, bucket, prefix, trainerutils.DefaultTrainingHubRuntime, + sdkInstallExports, + installKubeflowScript, torchrunFailureNotebookName, ) command := []string{"/bin/sh", "-c", shellCmd} diff --git a/tests/trainer/sdk_tests/fashion_mnist_tests.go b/tests/trainer/sdk_tests/fashion_mnist_tests.go index 08f3e066f..26936353f 100644 --- a/tests/trainer/sdk_tests/fashion_mnist_tests.go +++ b/tests/trainer/sdk_tests/fashion_mnist_tests.go @@ -35,10 +35,8 @@ import ( ) const ( - notebookName = "mnist.ipynb" - notebookPath = "resources/" + notebookName - installScriptPath = "resources/disconnected_env/install_kubeflow.py" - installKubeflowScript = "install_kubeflow.py" + notebookName = "mnist.ipynb" + notebookPath = "resources/" + notebookName ) // CPU Only - Distributed Training @@ -52,8 +50,8 @@ func RunFashionMnistCpuDistributedTraining(t *testing.T) { trainerutils.EnsureNotebookServiceAccount(t, test, namespace.Name) // RBACs setup - userName := common.GetNotebookUserName(test) - userToken := common.GenerateNotebookUserToken(test) + userToken := common.GetNotebookUserTokenFromEnv(test) + userName := common.GetNotebookUserNameFromEnv(test, userToken) support.CreateUserRoleBindingWithClusterRole(test, userName, namespace.Name, "admin") // ClusterRoleBinding for cluster-scoped resources (ClusterTrainingRuntimes) - minimal get/list/watch access trainerutils.CreateUserClusterRoleBindingForTrainerRuntimes(test, userName) @@ -91,6 +89,7 @@ func RunFashionMnistCpuDistributedTraining(t *testing.T) { support.StorageClassName(storageClass.Name), ) + sdkInstallExports := buildKubeflowInstallExports() shellCmd := fmt.Sprintf( "set -e; "+ "export OPENSHIFT_API_URL='%s'; export NOTEBOOK_TOKEN='%s'; "+ @@ -101,6 +100,7 @@ func RunFashionMnistCpuDistributedTraining(t *testing.T) { "export AWS_STORAGE_BUCKET_MNIST_DIR='%s'; "+ "export TRAINING_RUNTIME='%s'; "+ "export GPU_TYPE='cpu'; "+ + "%s"+ "python -m pip install --quiet --no-cache-dir ipykernel papermill boto3==1.34.162 && "+ "python /opt/app-root/notebooks/%s && "+ "if python -m papermill -k python3 /opt/app-root/notebooks/%s /opt/app-root/src/out.ipynb --log-output; "+ @@ -108,6 +108,7 @@ func RunFashionMnistCpuDistributedTraining(t *testing.T) { support.GetOpenShiftApiUrl(test), userToken, namespace.Name, rwxPvc.Name, endpoint, accessKey, secretKey, bucket, prefix, trainerutils.DefaultClusterTrainingRuntime, + sdkInstallExports, installKubeflowScript, notebookName, ) @@ -143,8 +144,8 @@ func RunFashionMnistKueueCpuDistributedTraining(t *testing.T) { trainerutils.EnsureNotebookServiceAccount(t, test, namespace.Name) // RBACs setup - userName := common.GetNotebookUserName(test) - userToken := common.GenerateNotebookUserToken(test) + userToken := common.GetNotebookUserTokenFromEnv(test) + userName := common.GetNotebookUserNameFromEnv(test, userToken) support.CreateUserRoleBindingWithClusterRole(test, userName, namespace.Name, "admin") // ClusterRoleBinding for cluster-scoped resources (ClusterTrainingRuntimes) - minimal get/list/watch access trainerutils.CreateUserClusterRoleBindingForTrainerRuntimes(test, userName) @@ -218,6 +219,7 @@ func RunFashionMnistKueueCpuDistributedTraining(t *testing.T) { support.StorageClassName(storageClass.Name), ) + sdkInstallExports := buildKubeflowInstallExports() shellCmd := fmt.Sprintf( "set -e; "+ "export OPENSHIFT_API_URL='%s'; export NOTEBOOK_TOKEN='%s'; "+ @@ -229,6 +231,7 @@ func RunFashionMnistKueueCpuDistributedTraining(t *testing.T) { "export TRAINING_RUNTIME='%s'; "+ "export GPU_TYPE='cpu'; "+ "export KUEUE_QUEUE_NAME='%s'; "+ + "%s"+ "python -m pip install --quiet --no-cache-dir ipykernel papermill boto3==1.34.162 && "+ "python /opt/app-root/notebooks/%s && "+ "if python -m papermill -k python3 /opt/app-root/notebooks/%s /opt/app-root/src/out.ipynb --log-output; "+ @@ -237,6 +240,7 @@ func RunFashionMnistKueueCpuDistributedTraining(t *testing.T) { endpoint, accessKey, secretKey, bucket, prefix, trainerutils.DefaultClusterTrainingRuntime, customLocalQueue.Name, + sdkInstallExports, installKubeflowScript, notebookName, ) diff --git a/tests/trainer/sdk_tests/lora_traininghub_tests.go b/tests/trainer/sdk_tests/lora_traininghub_tests.go index 8e3847322..e16c4414f 100644 --- a/tests/trainer/sdk_tests/lora_traininghub_tests.go +++ b/tests/trainer/sdk_tests/lora_traininghub_tests.go @@ -46,17 +46,22 @@ func RunLoraTrainingHubMultiGpuDistributedTraining(t *testing.T) { trainerutils.EnsureNotebookServiceAccount(t, test, namespace.Name) // RBACs setup - userName := common.GetNotebookUserName(test) - userToken := common.GenerateNotebookUserToken(test) + userToken := common.GetNotebookUserTokenFromEnv(test) + userName := common.GetNotebookUserNameFromEnv(test, userToken) support.CreateUserRoleBindingWithClusterRole(test, userName, namespace.Name, "admin") // ClusterRoleBinding for cluster-scoped resources (ClusterTrainingRuntimes) - minimal get/list/watch access trainerutils.CreateUserClusterRoleBindingForTrainerRuntimes(test, userName) - // Create ConfigMap with notebook + // Create ConfigMap with notebook and kubeflow install script localPath := loraNotebookPath nb, err := os.ReadFile(localPath) test.Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("failed to read notebook: %s", localPath)) - cm := support.CreateConfigMap(test, namespace.Name, map[string][]byte{loraNotebookName: nb}) + installScript, err := os.ReadFile(installScriptPath) + test.Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("failed to read install script: %s", installScriptPath)) + cm := support.CreateConfigMap(test, namespace.Name, map[string][]byte{ + loraNotebookName: nb, + installKubeflowScript: installScript, + }) // Build command with parameters and pinned deps, and print definitive status line to logs endpoint, endpointOK := support.GetStorageBucketDefaultEndpoint() @@ -82,6 +87,7 @@ func RunLoraTrainingHubMultiGpuDistributedTraining(t *testing.T) { support.StorageClassName(storageClass.Name), ) + sdkInstallExports := buildKubeflowInstallExports() shellCmd := fmt.Sprintf( "set -e; "+ "export OPENSHIFT_API_URL='%s'; export NOTEBOOK_USER_TOKEN='%s'; "+ @@ -92,12 +98,16 @@ func RunLoraTrainingHubMultiGpuDistributedTraining(t *testing.T) { "export AWS_STORAGE_BUCKET='%s'; "+ "export AWS_STORAGE_BUCKET_LORA_DIR='%s'; "+ "export TRAINING_RUNTIME='%s'; "+ + "%s"+ "python -m pip install --quiet --no-cache-dir --break-system-packages ipykernel papermill boto3==1.34.162 && "+ + "python /opt/app-root/notebooks/%s && "+ "if python -m papermill -k python3 /opt/app-root/notebooks/%s /opt/app-root/src/out.ipynb --log-output; "+ "then echo 'NOTEBOOK_STATUS: SUCCESS'; else echo 'NOTEBOOK_STATUS: FAILURE'; fi; sleep infinity", support.GetOpenShiftApiUrl(test), userToken, namespace.Name, rwxPvc.Name, endpoint, accessKey, secretKey, bucket, prefix, trainerutils.DefaultTrainingHubRuntime, + sdkInstallExports, + installKubeflowScript, loraNotebookName, ) command := []string{"/bin/sh", "-c", shellCmd} diff --git a/tests/trainer/sdk_tests/osft_traininghub_tests.go b/tests/trainer/sdk_tests/osft_traininghub_tests.go index 243ba3253..b4ca2d0a6 100644 --- a/tests/trainer/sdk_tests/osft_traininghub_tests.go +++ b/tests/trainer/sdk_tests/osft_traininghub_tests.go @@ -46,17 +46,22 @@ func RunOsftTrainingHubMultiGpuDistributedTraining(t *testing.T) { trainerutils.EnsureNotebookServiceAccount(t, test, namespace.Name) // RBACs setup - userName := common.GetNotebookUserName(test) - userToken := common.GenerateNotebookUserToken(test) + userToken := common.GetNotebookUserTokenFromEnv(test) + userName := common.GetNotebookUserNameFromEnv(test, userToken) support.CreateUserRoleBindingWithClusterRole(test, userName, namespace.Name, "admin") // ClusterRoleBinding for cluster-scoped resources (ClusterTrainingRuntimes) - minimal get/list/watch access trainerutils.CreateUserClusterRoleBindingForTrainerRuntimes(test, userName) - // Create ConfigMap with notebook + // Create ConfigMap with notebook and kubeflow install script localPath := osftNotebookPath nb, err := os.ReadFile(localPath) test.Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("failed to read notebook: %s", localPath)) - cm := support.CreateConfigMap(test, namespace.Name, map[string][]byte{osftNotebookName: nb}) + installScript, err := os.ReadFile(installScriptPath) + test.Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("failed to read install script: %s", installScriptPath)) + cm := support.CreateConfigMap(test, namespace.Name, map[string][]byte{ + osftNotebookName: nb, + installKubeflowScript: installScript, + }) // Build command with parameters and pinned deps, and print definitive status line to logs endpoint, endpointOK := support.GetStorageBucketDefaultEndpoint() @@ -82,6 +87,7 @@ func RunOsftTrainingHubMultiGpuDistributedTraining(t *testing.T) { support.StorageClassName(storageClass.Name), ) + sdkInstallExports := buildKubeflowInstallExports() shellCmd := fmt.Sprintf( "set -e; "+ "export OPENSHIFT_API_URL='%s'; export NOTEBOOK_USER_TOKEN='%s'; "+ @@ -92,12 +98,16 @@ func RunOsftTrainingHubMultiGpuDistributedTraining(t *testing.T) { "export AWS_STORAGE_BUCKET='%s'; "+ "export AWS_STORAGE_BUCKET_OSFT_DIR='%s'; "+ "export TRAINING_RUNTIME='%s'; "+ + "%s"+ "python -m pip install --quiet --no-cache-dir --break-system-packages ipykernel papermill boto3==1.34.162 && "+ + "python /opt/app-root/notebooks/%s && "+ "if python -m papermill -k python3 /opt/app-root/notebooks/%s /opt/app-root/src/out.ipynb --log-output; "+ "then echo 'NOTEBOOK_STATUS: SUCCESS'; else echo 'NOTEBOOK_STATUS: FAILURE'; fi; sleep infinity", support.GetOpenShiftApiUrl(test), userToken, namespace.Name, rwxPvc.Name, endpoint, accessKey, secretKey, bucket, prefix, trainerutils.DefaultTrainingHubRuntime, + sdkInstallExports, + installKubeflowScript, osftNotebookName, ) command := []string{"/bin/sh", "-c", shellCmd} diff --git a/tests/trainer/sdk_tests/rhai_features_tests.go b/tests/trainer/sdk_tests/rhai_features_tests.go index d1b5d977b..d29b6de4f 100644 --- a/tests/trainer/sdk_tests/rhai_features_tests.go +++ b/tests/trainer/sdk_tests/rhai_features_tests.go @@ -261,8 +261,8 @@ func runRhaiFeaturesTestWithConfig(t *testing.T, config RhaiFeatureConfig) { trainerutils.EnsureNotebookServiceAccount(t, test, namespace.Name) // RBACs setup for user (user token is used by notebook for Trainer API calls) - userName := common.GetNotebookUserName(test) - userToken := common.GenerateNotebookUserToken(test) + userToken := common.GetNotebookUserTokenFromEnv(test) + userName := common.GetNotebookUserNameFromEnv(test, userToken) CreateUserRoleBindingWithClusterRole(test, userName, namespace.Name, "admin") // ClusterRoleBinding for cluster-scoped resources (ClusterTrainingRuntimes) - minimal get/list/watch access trainerutils.CreateUserClusterRoleBindingForTrainerRuntimes(test, userName) @@ -379,11 +379,7 @@ func runRhaiFeaturesTestWithConfig(t *testing.T, config RhaiFeatureConfig) { secret := CreateSecret(test, namespace.Name, secretData) test.T().Logf("Created Data Connection secret: %s for cloud checkpoint storage", secret.Name) - dataConnectionExports = fmt.Sprintf( - "export DATA_CONNECTION_NAME='%s'; "+ - "export KUBEFLOW_INSTALL_FROM_GIT='true'; ", - secret.Name, - ) + dataConnectionExports = fmt.Sprintf("export DATA_CONNECTION_NAME='%s'; ", secret.Name) test.T().Logf("Data Connection configured for cloud checkpointing: %s", config.CheckpointOutputDir) } else if checkpointURI != nil { test.T().Logf("Warning: Cloud storage URI detected (%s) but Data Connection not created (credentials may be missing or unsupported scheme)", config.CheckpointOutputDir) @@ -401,6 +397,7 @@ func runRhaiFeaturesTestWithConfig(t *testing.T, config RhaiFeatureConfig) { // Build pip exports - GPU_TYPE tells install_kubeflow.py which Red Hat index to use pipExports := fmt.Sprintf("export GPU_TYPE='%s'; ", gpuType) + sdkInstallExports := buildKubeflowInstallExports() pipInstallFlags := "" // Set defaults for num_nodes and num_gpus_per_node if not specified @@ -431,6 +428,7 @@ func runRhaiFeaturesTestWithConfig(t *testing.T, config RhaiFeatureConfig) { "export NUM_GPUS_PER_NODE='%d'; "+ "%s"+ // S3 exports (if configured) "%s"+ // Data Connection exports (if configured) + "%s"+ // SDK install exports (git/version/index override) "%s"+ // PyPI/GPU_TYPE exports "python -m pip install --quiet --no-cache-dir %s papermill ipykernel boto3==1.34.162 && "+ "python /opt/app-root/notebooks/install_kubeflow.py && "+ @@ -449,6 +447,7 @@ func runRhaiFeaturesTestWithConfig(t *testing.T, config RhaiFeatureConfig) { numGpusPerNode, s3Exports, dataConnectionExports, + sdkInstallExports, pipExports, pipInstallFlags, config.NotebookName, diff --git a/tests/trainer/sdk_tests/sdk_env_helpers.go b/tests/trainer/sdk_tests/sdk_env_helpers.go new file mode 100644 index 000000000..5b9797819 --- /dev/null +++ b/tests/trainer/sdk_tests/sdk_env_helpers.go @@ -0,0 +1,54 @@ +package sdk_tests + +import ( + "os" + "strings" +) + +const ( + installScriptPath = "resources/disconnected_env/install_kubeflow.py" + installKubeflowScript = "install_kubeflow.py" +) + +// buildKubeflowInstallExports builds a shell-export prefix that is injected into +// notebook pod commands before running install_kubeflow.py. +// +// Why this exists: +// - SDK tests start from host-side go test processes, but kubeflow installation +// happens inside notebook containers. +// - Host environment variables are not automatically available in the notebook +// process, so we explicitly export the selected install vars. +// +// Selection precedence: +// 1. KUBEFLOW_GIT_URL: install from git (sets KUBEFLOW_INSTALL_FROM_GIT=true) +// 2. KUBEFLOW_REQUIRED_VERSION: install that version from index/default index +// 3. Neither set: installer script falls back to its internal defaults +// +// Index behavior: +// - If KUBEFLOW_SDK_INDEX_URL is set, we export it as KUBEFLOW_PYPI_INDEX_URL +// for install_kubeflow.py to use as the package index. +func buildKubeflowInstallExports() string { + gitURL := strings.TrimSpace(os.Getenv("KUBEFLOW_GIT_URL")) + + version := strings.TrimSpace(os.Getenv("KUBEFLOW_REQUIRED_VERSION")) + + indexURL := strings.TrimSpace(os.Getenv("KUBEFLOW_SDK_INDEX_URL")) + + var exports strings.Builder + if gitURL != "" { + exports.WriteString("export KUBEFLOW_INSTALL_FROM_GIT='true'; ") + exports.WriteString("export KUBEFLOW_GIT_URL=" + shellQuote(gitURL) + "; ") + } else if version != "" { + exports.WriteString("export KUBEFLOW_REQUIRED_VERSION=" + shellQuote(version) + "; ") + } + + if indexURL != "" { + exports.WriteString("export KUBEFLOW_PYPI_INDEX_URL=" + shellQuote(indexURL) + "; ") + } + return exports.String() +} + +func shellQuote(value string) string { + return "'" + strings.ReplaceAll(value, "'", "'\"'\"'") + "'" +} + diff --git a/tests/trainer/sdk_tests/sft_traininghub_tests.go b/tests/trainer/sdk_tests/sft_traininghub_tests.go index f348e4dab..152a0bfaf 100644 --- a/tests/trainer/sdk_tests/sft_traininghub_tests.go +++ b/tests/trainer/sdk_tests/sft_traininghub_tests.go @@ -46,17 +46,22 @@ func RunSftTrainingHubMultiGpuDistributedTraining(t *testing.T) { trainerutils.EnsureNotebookServiceAccount(t, test, namespace.Name) // RBACs setup - userName := common.GetNotebookUserName(test) - userToken := common.GenerateNotebookUserToken(test) + userToken := common.GetNotebookUserTokenFromEnv(test) + userName := common.GetNotebookUserNameFromEnv(test, userToken) support.CreateUserRoleBindingWithClusterRole(test, userName, namespace.Name, "admin") // ClusterRoleBinding for cluster-scoped resources (ClusterTrainingRuntimes) - minimal get/list/watch access trainerutils.CreateUserClusterRoleBindingForTrainerRuntimes(test, userName) - // Create ConfigMap with notebook + // Create ConfigMap with notebook and kubeflow install script localPath := sftNotebookPath nb, err := os.ReadFile(localPath) test.Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("failed to read notebook: %s", localPath)) - cm := support.CreateConfigMap(test, namespace.Name, map[string][]byte{sftNotebookName: nb}) + installScript, err := os.ReadFile(installScriptPath) + test.Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("failed to read install script: %s", installScriptPath)) + cm := support.CreateConfigMap(test, namespace.Name, map[string][]byte{ + sftNotebookName: nb, + installKubeflowScript: installScript, + }) // Build command with parameters and pinned deps, and print definitive status line to logs endpoint, endpointOK := support.GetStorageBucketDefaultEndpoint() @@ -82,6 +87,7 @@ func RunSftTrainingHubMultiGpuDistributedTraining(t *testing.T) { support.StorageClassName(storageClass.Name), ) + sdkInstallExports := buildKubeflowInstallExports() shellCmd := fmt.Sprintf( "set -e; "+ "export OPENSHIFT_API_URL='%s'; export NOTEBOOK_USER_TOKEN='%s'; "+ @@ -92,12 +98,16 @@ func RunSftTrainingHubMultiGpuDistributedTraining(t *testing.T) { "export AWS_STORAGE_BUCKET='%s'; "+ "export AWS_STORAGE_BUCKET_SFT_DIR='%s'; "+ "export TRAINING_RUNTIME='%s'; "+ + "%s"+ "python -m pip install --quiet --no-cache-dir --break-system-packages ipykernel papermill boto3==1.34.162 && "+ + "python /opt/app-root/notebooks/%s && "+ "if python -m papermill -k python3 /opt/app-root/notebooks/%s /opt/app-root/src/out.ipynb --log-output; "+ "then echo 'NOTEBOOK_STATUS: SUCCESS'; else echo 'NOTEBOOK_STATUS: FAILURE'; fi; sleep infinity", support.GetOpenShiftApiUrl(test), userToken, namespace.Name, rwxPvc.Name, endpoint, accessKey, secretKey, bucket, prefix, trainerutils.DefaultTrainingHubRuntime, + sdkInstallExports, + installKubeflowScript, sftNotebookName, ) command := []string{"/bin/sh", "-c", shellCmd}