| #!/usr/bin/env bash |
| |
| |
| set -uo pipefail |
|
|
| if [[ -n "${DLB_CONDA:-}" ]]; then |
| CONDA_BIN="${DLB_CONDA}" |
| elif command -v mamba >/dev/null 2>&1; then |
| CONDA_BIN="$(command -v mamba)" |
| elif command -v conda >/dev/null 2>&1; then |
| CONDA_BIN="$(command -v conda)" |
| else |
| echo "Neither mamba nor conda is available; set DLB_CONDA to an explicit executable." >&2 |
| exit 127 |
| fi |
|
|
| ALL_ENVIRONMENTS=( |
| dlb-flm dlb-langflow dlb-duo dlb-mdlm dlb-candi dlb-rdlm dlb-sdtt dlb-di4c dlb-eval |
| ) |
| if [[ -n "${DLB_ENV_NAMES:-}" ]]; then |
| IFS=',' read -r -a ENVIRONMENTS <<<"${DLB_ENV_NAMES}" |
| else |
| ENVIRONMENTS=("${ALL_ENVIRONMENTS[@]}") |
| fi |
|
|
| method_imports() { |
| case "$1" in |
| dlb-flm) printf '%s\n' 'datasets einops entmax flash_attn fsspec huggingface_hub hydra lightning numpy omegaconf pydantic requests rich scipy timm tokenizers torchmetrics tqdm transformers triton wandb yaml' ;; |
| dlb-langflow) printf '%s\n' 'datasets einops fsspec huggingface_hub pyarrow pydantic safetensors transformers yaml' ;; |
| dlb-duo) printf '%s\n' 'datasets einops flash_attn fsspec h5py huggingface_hub hydra lightning numpy omegaconf pydantic requests rich scipy timm tokenizers torchmetrics torchvision tqdm transformers triton wandb yaml' ;; |
| dlb-mdlm) printf '%s\n' 'causal_conv1d datasets einops eval_type_backport flash_attn fsspec huggingface_hub hydra lightning mamba_ssm numpy omegaconf pydantic requests rich timm tokenizers torchmetrics transformers wandb yaml' ;; |
| dlb-candi) printf '%s\n' 'datasets einops evaluate flash_attn fsspec huggingface_hub hydra lightning numpy omegaconf pydantic requests rich scipy tokenizers torchmetrics tqdm transformers yaml' ;; |
| dlb-rdlm) printf '%s\n' 'accelerate datasets einops eval_type_backport fsspec huggingface_hub hydra numpy omegaconf pydantic requests scipy tokenizers tqdm transformers wandb yaml' ;; |
| dlb-sdtt|dlb-di4c) printf '%s\n' 'datasets einops flash_attn fsspec huggingface_hub hydra lightning loguru mauve nltk numpy omegaconf pandas pydantic requests safetensors tensorboard timm tokenizers torchdata tqdm transformers wandb yaml' ;; |
| dlb-eval) printf '%s\n' 'accelerate datasets evaluate fsspec mauve pydantic sacrebleu scipy tokenizers transformers yaml' ;; |
| *) return 1 ;; |
| esac |
| } |
|
|
| emit_failure() { |
| python3 - "$1" "$2" "${3:-}" <<'PY' |
| import json |
| import sys |
|
|
| record = { |
| "environment": sys.argv[1], |
| "python": None, |
| "torch": None, |
| "torch_cuda": None, |
| "cuda_available": False, |
| "imports": {}, |
| "error": sys.argv[2], |
| } |
| if len(sys.argv) > 3 and sys.argv[3]: |
| record["diagnostic"] = sys.argv[3][-8000:] |
| print(json.dumps(record, sort_keys=True, allow_nan=False)) |
| PY |
| } |
|
|
| probe_diagnostic() { |
| local reason="$1" |
| local raw_output="$2" |
| local stderr_path="$3" |
| local stderr_output="" |
|
|
| [[ "${DLB_VERIFY_DEBUG:-}" == "1" ]] || return 0 |
|
|
| if [[ -s "${stderr_path}" ]]; then |
| stderr_output="$(<"${stderr_path}")" |
| fi |
|
|
| printf '%s\n' "${reason}" |
| if [[ -n "${stderr_output}" ]]; then |
| printf 'stderr:\n%s\n' "${stderr_output}" |
| fi |
| if [[ -n "${raw_output}" ]]; then |
| printf 'stdout:\n%s\n' "${raw_output}" |
| fi |
| } |
|
|
| validate_probe() { |
| python3 - "$1" "$2" "$3" <<'PY' |
| import json |
| import sys |
|
|
| marker = "DLB_ENV_PROBE_V1:" |
| environment, modules, raw = sys.argv[1:] |
| expected_modules = modules.split() |
| lines = raw.splitlines() |
| if len(lines) != 1 or not lines[0].startswith(marker): |
| raise SystemExit(2) |
|
|
|
|
| def reject_constant(value): |
| raise ValueError(f"non-finite JSON constant: {value}") |
|
|
|
|
| def reject_duplicate_keys(pairs): |
| result = {} |
| for key, value in pairs: |
| if key in result: |
| raise ValueError(f"duplicate JSON key: {key}") |
| result[key] = value |
| return result |
|
|
|
|
| try: |
| record = json.loads( |
| lines[0][len(marker):], |
| parse_constant=reject_constant, |
| object_pairs_hook=reject_duplicate_keys, |
| ) |
| except (json.JSONDecodeError, ValueError): |
| raise SystemExit(2) |
| if not isinstance(record, dict) or record.get("environment") != environment: |
| raise SystemExit(2) |
| required = {"environment", "python", "torch", "torch_cuda", "cuda_available", "imports"} |
| optional = { |
| "torch_error", |
| "cuda_error", |
| "import_errors", |
| "runtime_checks", |
| "runtime_errors", |
| } |
| if not required <= record.keys() or not set(record) <= required | optional: |
| raise SystemExit(2) |
| if type(record["python"]) is not str: |
| raise SystemExit(2) |
| if not (type(record["torch"]) is str or record["torch"] is None): |
| raise SystemExit(2) |
| if not (type(record["torch_cuda"]) is str or record["torch_cuda"] is None): |
| raise SystemExit(2) |
| if type(record["cuda_available"]) is not bool or type(record["imports"]) is not dict: |
| raise SystemExit(2) |
| if set(record["imports"]) != set(expected_modules): |
| raise SystemExit(2) |
| if not all(type(value) is bool for value in record["imports"].values()): |
| raise SystemExit(2) |
| if "torch_error" in record and ( |
| type(record["torch_error"]) is not str or record["torch"] is not None |
| ): |
| raise SystemExit(2) |
| if "cuda_error" in record and ( |
| type(record["cuda_error"]) is not str or record["cuda_available"] |
| ): |
| raise SystemExit(2) |
| if "import_errors" in record: |
| errors = record["import_errors"] |
| failed_modules = {name for name, available in record["imports"].items() if not available} |
| if ( |
| type(errors) is not dict |
| or set(errors) != failed_modules |
| or not all(type(value) is str for value in errors.values()) |
| ): |
| raise SystemExit(2) |
| elif not all(record["imports"].values()): |
| raise SystemExit(2) |
| if "runtime_checks" in record: |
| checks = record["runtime_checks"] |
| if ( |
| type(checks) is not dict |
| or not all(type(key) is str and type(value) is bool for key, value in checks.items()) |
| ): |
| raise SystemExit(2) |
| failed_checks = {name for name, passed in checks.items() if not passed} |
| if "runtime_errors" in record: |
| errors = record["runtime_errors"] |
| if ( |
| type(errors) is not dict |
| or set(errors) != failed_checks |
| or not all(type(value) is str for value in errors.values()) |
| ): |
| raise SystemExit(2) |
| elif failed_checks: |
| raise SystemExit(2) |
| elif "runtime_errors" in record: |
| raise SystemExit(2) |
| print(json.dumps(record, sort_keys=True, allow_nan=False)) |
| healthy = ( |
| record["cuda_available"] |
| and all(record["imports"].values()) |
| and all(record.get("runtime_checks", {}).values()) |
| and record["torch"] is not None |
| ) |
| raise SystemExit(0 if healthy else 1) |
| PY |
| } |
|
|
| failures=() |
| for environment in "${ENVIRONMENTS[@]}"; do |
| if ! imports="$(method_imports "${environment}")"; then |
| emit_failure "${environment}" "unknown environment" |
| failures+=("${environment}") |
| continue |
| fi |
|
|
| if ! probe_stderr="$(mktemp "${TMPDIR:-/tmp}/dlb-env-probe.XXXXXX")"; then |
| emit_failure "${environment}" "verification probe failed" |
| failures+=("${environment}") |
| continue |
| fi |
| if ! probe_script="$(mktemp "${TMPDIR:-/tmp}/dlb-env-probe.XXXXXX")"; then |
| emit_failure "${environment}" "verification probe failed" |
| rm -f -- "${probe_stderr}" |
| failures+=("${environment}") |
| continue |
| fi |
|
|
| if ! cat >"${probe_script}" <<'PY' |
| import importlib |
| import json |
| import platform |
| import sys |
|
|
| marker = "DLB_ENV_PROBE_V1:" |
| environment = sys.argv[1] |
| modules = sys.argv[2:] |
| record = { |
| "environment": environment, |
| "python": platform.python_version(), |
| "torch": None, |
| "torch_cuda": None, |
| "cuda_available": False, |
| "imports": {}, |
| } |
| torch_module = None |
| try: |
| import torch |
|
|
| torch_module = torch |
| record["torch"] = torch.__version__ |
| record["torch_cuda"] = torch.version.cuda |
| record["cuda_available"] = torch.cuda.is_available() |
| if not record["cuda_available"]: |
| record["cuda_error"] = "CUDA is unavailable" |
| except Exception as error: |
| record["torch_error"] = str(error) |
|
|
| for module in modules: |
| try: |
| importlib.import_module(module) |
| record["imports"][module] = True |
| except Exception as error: |
| record["imports"][module] = False |
| record.setdefault("import_errors", {})[module] = str(error) |
|
|
| def version_tuple(value): |
| parts = [] |
| for raw in value.replace("-", ".").split("."): |
| if not raw.isdigit(): |
| break |
| parts.append(int(raw)) |
| return tuple(parts) |
|
|
| runtime_checks = {} |
| runtime_errors = {} |
| if record["imports"].get("flash_attn") and torch_module is not None: |
| try: |
| import importlib.metadata |
|
|
| flash_version = importlib.metadata.version("flash-attn") |
| except Exception: |
| flash_version = "" |
| if flash_version and version_tuple(flash_version) >= (2, 8, 0): |
| check_name = "flash_attn_torch_wrap_triton" |
| passed = hasattr(torch_module.library, "wrap_triton") |
| runtime_checks[check_name] = passed |
| if not passed: |
| runtime_errors[check_name] = ( |
| f"flash-attn {flash_version} requires torch.library.wrap_triton " |
| f"but torch {record['torch']} does not provide it" |
| ) |
| if runtime_checks: |
| record["runtime_checks"] = runtime_checks |
| if runtime_errors: |
| record["runtime_errors"] = runtime_errors |
|
|
| print(marker + json.dumps(record, sort_keys=True, allow_nan=False)) |
| raise SystemExit(0) |
| PY |
| then |
| emit_failure "${environment}" "verification probe failed" |
| rm -f -- "${probe_stderr}" "${probe_script}" |
| failures+=("${environment}") |
| continue |
| fi |
|
|
| probe_output="" |
| manager_status=0 |
| probe_output="$("${CONDA_BIN}" run -n "${environment}" python \ |
| "${probe_script}" "${environment}" ${imports} 2>"${probe_stderr}")" || manager_status=$? |
| diagnostic="" |
| if ((manager_status != 0)); then |
| diagnostic="$( |
| probe_diagnostic "conda run exited with status ${manager_status}" \ |
| "${probe_output}" "${probe_stderr}" |
| )" |
| fi |
|
|
| if ((manager_status != 0)); then |
| emit_failure "${environment}" "verification probe failed" "${diagnostic}" |
| rm -f -- "${probe_stderr}" "${probe_script}" |
| failures+=("${environment}") |
| continue |
| fi |
|
|
| probe_status=0 |
| record="$(validate_probe "${environment}" "${imports}" "${probe_output}")" || probe_status=$? |
| if ((probe_status <= 1)); then |
| printf '%s\n' "${record}" |
| if ((probe_status != 0)); then |
| failures+=("${environment}") |
| fi |
| else |
| diagnostic="$( |
| probe_diagnostic "probe output was not valid DLB_ENV_PROBE_V1 JSON" \ |
| "${probe_output}" "${probe_stderr}" |
| )" |
| emit_failure "${environment}" "verification probe failed" "${diagnostic}" |
| failures+=("${environment}") |
| fi |
| rm -f -- "${probe_stderr}" "${probe_script}" |
| done |
|
|
| if ((${#failures[@]})); then |
| echo "Environment verification failed: ${failures[*]}" >&2 |
| exit 1 |
| fi |
|
|
