Allow a custom pyproject.toml template (#75)

* Allow a custom pyproject.toml template

This update enables users to provide a custom pyproject.toml file as
a template for the files generated by seed-env.

Here is how the new functionality works:
* The tool can now accept a user-provided pyproject.toml template by
  using a `--template-pyproject-toml` flag.
* By default, if a pyproject.toml file is present at the project's
  root when the seed-env CLI is called, the tool will automatically
  use it as the template.
* The template file used for input cannot be the same as the
  pyproject.toml file generated in the --output-dir.
* Disallow a pre-existing pyproject.toml in the --output-dir.

* Fix ruff format

Author: kanglant

python_seed_env/src/seed_env/cli.py

@@ -83,6 +83,12 @@ def main():
   )
 
   # --- Common Arguments ---
+  parser.add_argument(
+    "--template-pyproject-toml",
+    type=str,
+    default=None,
+    help="Path to a custom pyproject.toml file to use as a template.",
+  )
   parser.add_argument(
     "--seed-config",
     type=str,
@@ -207,6 +213,7 @@ def main():
       hardware=args.hardware,
       build_pypi_package=args.build_pypi_package,
       output_dir=args.output_dir,
+      template_pyproject_toml=args.template_pyproject_toml,
     )
     # Core function
     host_env_seeder.seed_environment()

python_seed_env/src/seed_env/config.py

@@ -52,12 +52,9 @@
   "nvidia-nvvm",
   "jax-cuda12-plugin",
   "jax-cuda13-plugin",
-  # "jax-cuda12-plugin[with-cuda]",
   "jax-cuda12-pjrt",
   "jax-cuda13-pjrt",
   "transformer-engine",
-  # "transformer-engine[jax]",
-  # "transformer-engine[pytorch]",
 ]
 
 TPU_SPECIFIC_DEPS = [

python_seed_env/src/seed_env/core.py

@@ -16,15 +16,18 @@
 
 import os
 import logging
+import shutil
 import yaml
 from importlib.resources import files
 from seed_env.seeder import Seeder
 from seed_env.utils import generate_minimal_pyproject_toml
 from seed_env.git_utils import download_remote_git_file
 from seed_env.uv_utils import (
+  set_exact_python_requirement_in_project_toml,
   build_seed_env,
   build_pypi_package,
   merge_project_toml_files,
+  replace_dependencies_in_project_toml,
 )
 
 logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
@@ -59,13 +62,15 @@ def __init__(
     hardware: str,
     build_pypi_package: bool,
     output_dir: str,
+    template_pyproject_toml: str = None,
   ):
     self.host_name = host_name
     self.host_source_type = host_source_type
     self.host_github_org_repo = host_github_org_repo
     self.host_requirements_file_path = host_requirements_file_path
     self.host_commit = host_commit
     self.seed_config_input = seed_config
+    self.template_pyproject_toml = template_pyproject_toml
     self.loaded_seed_config = None
     self.seed_tag_or_commit = seed_tag_or_commit
     self.python_versions = python_version.split(",")
@@ -143,6 +148,31 @@ def seed_environment(self):
     os.makedirs(self.output_dir, exist_ok=True)
     self.output_dir = os.path.abspath(self.output_dir)
 
+    # Determine the template for pyproject.toml. The explicit CLI argument takes precedence.
+    template_path = self.template_pyproject_toml
+    if not template_path and os.path.isfile("./pyproject.toml"):
+      template_path = os.path.abspath("./pyproject.toml")
+      logging.info(
+        f"Found pyproject.toml in the current directory. Using it as a template: {template_path}"
+      )
+
+    # Pre-flight check: Ensure the output directory root is clean of a pyproject.toml, as we will generate one.
+    final_pyproject_path = os.path.join(self.output_dir, "pyproject.toml")
+    if os.path.isfile(final_pyproject_path):
+      # Check for the specific edge case where the output directory is the project root
+      # and the existing pyproject.toml is the one we are using as a template.
+      if template_path and os.path.samefile(template_path, final_pyproject_path):
+        raise FileExistsError(
+          f"The output directory ('{self.output_dir}') contains a 'pyproject.toml', which was found to be used as a template. "
+          "Running this would overwrite the original template file. Please use a different --output-dir; or move the"
+          "existing pyproject.toml to a different location and use the --template-pyproject-toml flag to specify its new location."
+        )
+      # General case: the output directory contains a pre-existing pyproject.toml.
+      raise FileExistsError(
+        f"A pyproject.toml file already exists in the output directory: {self.output_dir}. "
+        "Please provide a clean directory or remove the file to avoid accidentaly overwriting it."
+      )
+
     # Create a directory for storing the downloaded requirements file
     self.download_dir = "downloaded_base_and_seed_requirements"
     os.makedirs(self.download_dir, exist_ok=True)
@@ -179,30 +209,36 @@ def seed_environment(self):
       f"Using {self.seeder.pypi_project_name} at tag/commit {self.seed_tag_or_commit} on {self.seeder.github_org_repo} as seed"
     )
 
-    # Remove pyproject.toml if it exists, as we will generate a new one with merge_project_toml_files
-    pyproject_file = os.path.join(self.output_dir, "pyproject.toml")
-    if os.path.isfile(pyproject_file):
-      os.remove(pyproject_file)
-      logging.info(f"Removed existing pyproject.toml file: {pyproject_file}")
-
     versioned_project_toml_files = []
     for python_version in self.python_versions:
       # Generate a subdir for each python version
       versioned_output_dir = (
         self.output_dir + "/python" + python_version.replace(".", "_")
       )
-      versioned_project_toml_files.append(versioned_output_dir + "/pyproject.toml")
       os.makedirs(versioned_output_dir, exist_ok=True)
+      versioned_pyproject_path = os.path.join(versioned_output_dir, "pyproject.toml")
+      versioned_project_toml_files.append(versioned_pyproject_path)
 
       # 3. Download the seed lock file for the specified Python version
       SEED_LOCK_FILE = os.path.abspath(
         self.seeder.download_seed_lock_requirement(python_version)
       )
 
-      # 4. Generate a minimal pyproject.toml file for the specified Python version to the output directory
-      generate_minimal_pyproject_toml(
-        self.host_name, python_version, versioned_output_dir
-      )
+      # 4. Generate a pyproject.toml file for the specified Python version.
+      if template_path:
+        logging.info(f"Using template {template_path} for Python {python_version}")
+        shutil.copy(template_path, versioned_pyproject_path)
+        # Clear any existing dependencies from the template to start fresh.
+        replace_dependencies_in_project_toml([], versioned_pyproject_path)
+        # Update the python version in the copied template to be specific for this build pass.
+        set_exact_python_requirement_in_project_toml(
+          python_version, versioned_pyproject_path
+        )
+      else:
+        logging.info(f"Generating minimal pyproject.toml for Python {python_version}")
+        generate_minimal_pyproject_toml(
+          self.host_name, python_version, versioned_output_dir
+        )
 
       # Construct the host lock file name
       HOST_LOCK_FILE_NAME = f"{self.host_name.replace('-', '_')}_requirements_lock_{python_version.replace('.', '_')}.txt"
@@ -220,6 +256,12 @@ def seed_environment(self):
     merge_project_toml_files(versioned_project_toml_files, self.output_dir)
 
     # 6. Build pypi package
+    # TODO(kanglant): Assume where the seed-env cli is called is the project root
+    # and move the generated pyproject.toml to the project root? If there is an old
+    # pyproject.toml here, this behavior will overwrite it. I think this is risky.
+    # Another option is to copy the source files, specified in the new pyproject.toml,
+    # from the project root to the output_dir folder. Then perform a fully isolated
+    # build at the output_dir.
     if self.build_pypi_package:
       # Use the new pyproject.toml file at the output dir to build the package.
       build_pypi_package(self.output_dir)

python_seed_env/src/seed_env/utils.py

@@ -98,21 +98,15 @@ def generate_minimal_pyproject_toml(
 
 [project]
 name = "{project_name}"
-description = "{project_name} is a simple, performant and scalable Jax LLM!"
 version = "0.0.1"
-readme = "README.md"
 license = "Apache-2.0"
-license-files = ["LICENSE"]
-keywords = ["llm", "jax", "llama", "mistral", "mixtral", "gemma", "deepseek"]
 requires-python = "=={python_version}.*"
 dependencies = [
 ]
 classifiers = [
-    "Development Status :: 4 - Beta",
     "Programming Language :: Python",
 ]
 
-# TODO(kanglant): Remove this once maxtext src-layout restructure done.
 [tool.hatch.build.targets.wheel]
 packages = ["{project_name}"]
 

python_seed_env/src/seed_env/uv_utils.py

@@ -83,8 +83,6 @@ def build_seed_env(
     output_dir,
     "-r",
     seed_lock_file,
-    "--extra-index-url",
-    "https://pypi.nvidia.com",
   ]
   run_command(command)
 
@@ -103,8 +101,6 @@ def build_seed_env(
   ]
   run_command(command)
 
-  _remove_hardware_specific_deps(hardware, pyproject_file, output_dir)
-
   command = [
     "uv",
     "export",
@@ -241,7 +237,10 @@ def replace_dependencies_in_project_toml(new_deps_list: list, filepath: str):
   This function reads the specified pyproject.toml file, finds the existing project dependencies array,
   and replaces it with the provided new_deps_list list. The updated content is then written back to the file.
   """
-  new_deps = 'dependencies = [\n    "' + '",\n    "'.join(new_deps_list) + '",\n]'
+  if new_deps_list:
+    new_deps = 'dependencies = [\n    "' + '",\n    "'.join(new_deps_list) + '",\n]'
+  else:
+    new_deps = "dependencies = []"
 
   dependencies_regex = re.compile(
     r"^dependencies\s*=\s*\[(\n+\s*.*,\s*)*[\n\r]*\]", re.MULTILINE
@@ -277,6 +276,37 @@ def replace_python_requirement_in_project_toml(min_python: str, filepath: str):
     f.write(new_content)
 
 
+def set_exact_python_requirement_in_project_toml(python_version: str, filepath: str):
+  """
+  Sets or adds the requires-python section in a pyproject.toml file to an exact version series.
+
+  Args:
+      python_version (str): The target Python version (e.g., '3.12').
+      filepath (str): Path to the pyproject.toml file to update.
+  """
+  python_req_regex = re.compile(r'requires-python\s*=\s*".*?"')
+  project_header_regex = re.compile(r"^\[project\]", re.MULTILINE)
+  new_requires_line = f'requires-python = "=={python_version}.*"'
+
+  with open(filepath, "r", encoding="utf-8") as f:
+    content = f.read()
+
+  if python_req_regex.search(content):
+    # If 'requires-python' exists, substitute it.
+    new_content = python_req_regex.sub(new_requires_line, content)
+  elif project_header_regex.search(content):
+    # If it doesn't exist but [project] does, add it after the [project] header.
+    new_content = project_header_regex.sub(
+      f"[project]\n{new_requires_line}", content, count=1
+    )
+  else:
+    logging.error("No project table found in the template pyproject.toml.")
+    raise
+
+  with open(filepath, "w", encoding="utf-8") as f:
+    f.write(new_content)
+
+
 def lock_to_lower_bound_project(host_lock_file: str, pyproject_toml: str):
   """
   Updates the dependencies in a pyproject.toml file to use lower-bound versions based on a lock file.

python_seed_env/tests/test_core.py

@@ -57,10 +57,8 @@ def test_environment_seeder_init_invalid_seed():
 def test_seed_environment_remote(mocker, tmp_path):
   # Mock all external dependencies
   mock_download = mocker.patch(
-    "seed_env.core.download_remote_git_file", return_value=str(tmp_path / "host.txt")
-  )
-  mock_generate_pyproject = mocker.patch(
-    "seed_env.core.generate_minimal_pyproject_toml"
+    "seed_env.core.download_remote_git_file",
+    return_value=str(tmp_path / "host.txt"),
   )
   mock_merge_project_toml_files = mocker.patch("seed_env.core.merge_project_toml_files")
   mock_build_env = mocker.patch("seed_env.core.build_seed_env")
@@ -74,6 +72,11 @@ def test_seed_environment_remote(mocker, tmp_path):
   )
   mocker.patch("seed_env.core.Seeder", return_value=mock_seeder_instance)
 
+  # 4. Instantiate and run the seeder.
+  template_toml_path = tmp_path / "pyproject.toml"
+  template_toml_path.write_text(
+    '[project]\nname = "myproj"\nreadme = "README.md"\n[tool.hatch.build.targets.wheel]\npackages = ["myproj"]'
+  )
   seeder = EnvironmentSeeder(
     host_name="myproj",
     host_source_type="remote",
@@ -86,16 +89,17 @@ def test_seed_environment_remote(mocker, tmp_path):
     hardware="cpu",
     build_pypi_package=True,
     output_dir=str(tmp_path / "output"),
+    template_pyproject_toml=str(template_toml_path),
   )
   seeder.seed_environment()
 
   # Assert all mocks were called
   assert mock_download.called
-  assert mock_generate_pyproject.called
+  # assert mock_generate_pyproject.called
   assert mock_build_env.called
   assert mock_merge_project_toml_files.called
   assert mock_build_pypi.called
-  assert mock_seeder_instance.download_seed_lock_requirement.called
+  mock_seeder_instance.download_seed_lock_requirement.assert_called_with("3.12")
 
 
 def test_seed_environment_local_file_not_found(mocker, tmp_path):

python_seed_env/tests/test_utils.py

@@ -159,7 +159,7 @@ def test_build_seed_env_calls_run_command(mocker, tmp_path):
   assert mock_os_remove.called
   assert mock_lock_to_lower_bound_project.called
   # Check for the expected uv remove command
-  assert mock_remove_hardware_specific_deps.call_count == 2
+  assert mock_remove_hardware_specific_deps.call_count == 1
 
   # Collect all commands passed to run_command
   commands = [call.args[0] for call in mock_run_command.call_args_list]