From 1461d1436a0135f578d0d6dbca1bf544dc54e668 Mon Sep 17 00:00:00 2001 From: Kanglan Tang Date: Thu, 16 Oct 2025 01:18:54 +0000 Subject: [PATCH] Add --hardware=cpu flag to seed-env --- python_seed_env/src/seed_env/config.py | 2 +- python_seed_env/src/seed_env/uv_utils.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/python_seed_env/src/seed_env/config.py b/python_seed_env/src/seed_env/config.py index 8d9107c..fb703b1 100644 --- a/python_seed_env/src/seed_env/config.py +++ b/python_seed_env/src/seed_env/config.py @@ -20,7 +20,7 @@ DEFAULT_PYTHON_VERSION = "3.12" DEFAULT_HARDWARE = "tpu" DEFAULT_BUILD_PROJECT = False -SUPPORTED_HARDWARE = ["tpu", "gpu", "cuda12", "cuda13"] +SUPPORTED_HARDWARE = ["tpu", "gpu", "cuda12", "cuda13", "cpu"] TENSORFLOW_DEPS = [ "tensorflow", diff --git a/python_seed_env/src/seed_env/uv_utils.py b/python_seed_env/src/seed_env/uv_utils.py index 913d1b9..0e17059 100644 --- a/python_seed_env/src/seed_env/uv_utils.py +++ b/python_seed_env/src/seed_env/uv_utils.py @@ -381,6 +381,11 @@ def _remove_hardware_specific_deps(hardware: str, pyproject_file: str, output_di hardware_specific_deps_list = TPU_SPECIFIC_DEPS.copy() hardware_specific_deps_list.extend(CUDA12_SPECIFIC_DEPS) hardware_specific_deps_list.extend(TENSORFLOW_DEPS) + elif hardware == "cpu": + hardware_specific_deps_list = TPU_SPECIFIC_DEPS.copy() + hardware_specific_deps_list.extend(CUDA12_SPECIFIC_DEPS) + hardware_specific_deps_list.extend(CUDA13_SPECIFIC_DEPS) + hardware_specific_deps_list.extend(TENSORFLOW_DEPS) else: logging.warning(f"Unknown hardware {hardware}. Please use tpu or gpu.") return