Cleanup init and init_directory usage (#1368)

Author: nnarayen

truss/cli/cli.py

@@ -205,14 +205,14 @@ def image():
 )
 @click.option("-n", "--name", type=click.STRING)
 @click.option(
-    "--python-configuration/--no-python-configuration",
+    "--python-config/--no-python-config",
     type=bool,
     default=False,
     help="Uses the code first tooling to build models.",
 )
 @log_level_option
 @error_handling
-def init(target_directory, backend, name, python_configuration) -> None:
+def init(target_directory, backend, name, python_config) -> None:
     """Create a new truss.
 
     TARGET_DIRECTORY: A Truss is created in this directory
@@ -232,7 +232,7 @@ def init(target_directory, backend, name, python_configuration) -> None:
         target_directory=target_directory,
         build_config=build_config,
         model_name=model_name,
-        python_configuration=python_configuration,
+        python_config=python_config,
     )
     click.echo(f"Truss {model_name} was created in {tr_path.absolute()}")
 

truss/tests/conftest.py

@@ -20,7 +20,7 @@
     ServingImageBuilderContext,
 )
 from truss.contexts.local_loader.docker_build_emulator import DockerBuildEmulator
-from truss.truss_handle.build import init
+from truss.truss_handle.build import init_directory
 from truss.truss_handle.truss_handle import TrussHandle
 
 CUSTOM_MODEL_CODE = """
@@ -435,17 +435,15 @@ def dynamic_config_mount_dir(tmp_path, monkeypatch: pytest.MonkeyPatch):
 @pytest.fixture
 def custom_model_truss_dir_with_pre_and_post_no_example(tmp_path):
     dir_path = tmp_path / "custom_truss_with_pre_post_no_example"
-    handle = init(str(dir_path))
-    handle.spec.model_class_filepath.write_text(
-        CUSTOM_MODEL_CODE_WITH_PRE_AND_POST_PROCESS
-    )
+    th = TrussHandle(init_directory(dir_path))
+    th.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_WITH_PRE_AND_POST_PROCESS)
     yield dir_path
 
 
 @pytest.fixture
 def custom_model_truss_dir_with_hidden_files(tmp_path):
     truss_dir_path: Path = tmp_path / "custom_model_truss_dir_with_hidden_files"
-    _ = init(str(truss_dir_path))
+    init_directory(truss_dir_path)
     (truss_dir_path / "__pycache__").mkdir(parents=True, exist_ok=True)
     (truss_dir_path / ".git").mkdir(parents=True, exist_ok=True)
     (truss_dir_path / "__pycache__" / "test.cpython-311.pyc").touch()
@@ -458,7 +456,7 @@ def custom_model_truss_dir_with_hidden_files(tmp_path):
 @pytest.fixture
 def custom_model_truss_dir_with_truss_ignore(tmp_path):
     truss_dir_path: Path = tmp_path / "custom_model_truss_dir_with_truss_ignore"
-    _ = init(str(truss_dir_path))
+    init_directory(truss_dir_path)
     (truss_dir_path / "random_folder_1").mkdir(parents=True, exist_ok=True)
     (truss_dir_path / "random_folder_2").mkdir(parents=True, exist_ok=True)
     (truss_dir_path / "random_file_1.txt").touch()
@@ -478,19 +476,17 @@ def custom_model_truss_dir_with_truss_ignore(tmp_path):
 @pytest.fixture
 def custom_model_truss_dir_with_pre_and_post(tmp_path):
     dir_path = tmp_path / "custom_truss_with_pre_post"
-    handle = init(str(dir_path))
-    handle.spec.model_class_filepath.write_text(
-        CUSTOM_MODEL_CODE_WITH_PRE_AND_POST_PROCESS
-    )
-    handle.update_examples([Example("example1", {"inputs": [[0]]})])
+    th = TrussHandle(init_directory(dir_path))
+    th.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_WITH_PRE_AND_POST_PROCESS)
+    th.update_examples([Example("example1", {"inputs": [[0]]})])
     yield dir_path
 
 
 @pytest.fixture
 def custom_model_truss_dir_with_bundled_packages(tmp_path):
     truss_dir_path: Path = tmp_path / "custom_model_truss_dir_with_bundled_packages"
-    handle = init(str(truss_dir_path))
-    handle.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_USING_BUNDLED_PACKAGE)
+    th = TrussHandle(init_directory(truss_dir_path))
+    th.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_USING_BUNDLED_PACKAGE)
     packages_path = truss_dir_path / DEFAULT_BUNDLED_PACKAGES_DIR / "test_package"
     packages_path.mkdir(parents=True)
     with (packages_path / "test.py").open("w") as file:
@@ -501,11 +497,9 @@ def custom_model_truss_dir_with_bundled_packages(tmp_path):
 @pytest.fixture
 def custom_model_truss_dir_with_pre_and_post_str_example(tmp_path):
     dir_path = tmp_path / "custom_truss_with_pre_post_str_example"
-    handle = init(str(dir_path))
-    handle.spec.model_class_filepath.write_text(
-        CUSTOM_MODEL_CODE_WITH_PRE_AND_POST_PROCESS
-    )
-    handle.update_examples(
+    th = TrussHandle(init_directory(dir_path))
+    th.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_WITH_PRE_AND_POST_PROCESS)
+    th.update_examples(
         [
             Example(
                 "example1",
@@ -525,29 +519,27 @@ def custom_model_truss_dir_with_pre_and_post_str_example(tmp_path):
 @pytest.fixture
 def custom_model_truss_dir_with_pre_and_post_description(tmp_path):
     dir_path = tmp_path / "custom_truss_with_pre_post"
-    handle = init(str(dir_path))
-    handle.spec.model_class_filepath.write_text(
-        CUSTOM_MODEL_CODE_WITH_PRE_AND_POST_PROCESS
-    )
-    handle.update_description("This model adds 3 to all inputs")
+    th = TrussHandle(init_directory(dir_path))
+    th.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_WITH_PRE_AND_POST_PROCESS)
+    th.update_description("This model adds 3 to all inputs")
     yield dir_path
 
 
 @pytest.fixture
 def custom_model_truss_dir_for_gpu(tmp_path):
     dir_path = tmp_path / "custom_truss"
-    handle = init(str(dir_path))
-    handle.enable_gpu()
-    handle.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_FOR_GPU_TESTING)
+    th = TrussHandle(init_directory(dir_path))
+    th.enable_gpu()
+    th.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_FOR_GPU_TESTING)
     yield dir_path
 
 
 @pytest.fixture
 def custom_model_truss_dir_for_secrets(tmp_path):
     dir_path = tmp_path / "custom_truss"
-    handle = init(str(dir_path))
-    handle.add_secret("secret_name", "default_secret_value")
-    handle.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_FOR_SECRETS_TESTING)
+    th = TrussHandle(init_directory(dir_path))
+    th.add_secret("secret_name", "default_secret_value")
+    th.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_FOR_SECRETS_TESTING)
     yield dir_path
 
 
@@ -624,10 +616,10 @@ def _custom_model_from_code(
     where_dir: Path, truss_name: str, model_code: str, handle_ops: callable = None
 ) -> Path:
     dir_path = where_dir / truss_name
-    handle = init(str(dir_path))
+    th = TrussHandle(init_directory(dir_path))
     if handle_ops is not None:
-        handle_ops(handle)
-    handle.spec.model_class_filepath.write_text(model_code)
+        handle_ops(th)
+    th.spec.model_class_filepath.write_text(model_code)
     return dir_path
 
 

truss/tests/test_build.py

@@ -1,67 +1,18 @@
-from pathlib import Path
-
 from truss.base.truss_spec import TrussSpec
-from truss.truss_handle.build import init, init_directory, load
+from truss.truss_handle.build import init_directory, load
 from truss_chains.deployment import code_gen
 
 
 def test_truss_init(tmp_path):
-    dir_name = str(tmp_path)
-    init(dir_name)
-    spec = TrussSpec(Path(dir_name))
+    spec = TrussSpec(init_directory(tmp_path))
     assert spec.model_module_dir.exists()
     assert spec.data_dir.exists()
     assert spec.truss_dir == tmp_path
     assert spec.config_path.exists()
 
 
-def test_truss_init_with_data_file_and_requirements_file_and_bundled_packages(tmp_path):
-    dir_path = tmp_path / "truss"
-    dir_name = str(dir_path)
-
-    # Init data files
-    data_path = tmp_path / "data.txt"
-    with data_path.open("w") as data_file:
-        data_file.write("test")
-
-    # Init requirements file
-    req_file_path = tmp_path / "requirements.txt"
-    requirements = ["tensorflow==2.3.1", "uvicorn==0.12.2"]
-    with req_file_path.open("w") as req_file:
-        for req in requirements:
-            req_file.write(f"{req}\n")
-
-    # init bundled packages
-    packages_path = tmp_path / "dep_pkg"
-    packages_path.mkdir()
-    packages_path_file_py = packages_path / "file.py"
-    packages_path_init_py = packages_path / "__init__.py"
-    pkg_files = [packages_path_init_py, packages_path_file_py]
-    for pkg_file in pkg_files:
-        with pkg_file.open("w") as fh:
-            fh.write("test")
-
-    init(
-        dir_name,
-        data_files=[str(data_path)],
-        requirements_file=str(req_file_path),
-        bundled_packages=[str(packages_path)],
-    )
-    spec = TrussSpec(Path(dir_name))
-    assert spec.model_module_dir.exists()
-    assert spec.truss_dir == dir_path
-    assert spec.config_path.exists()
-    assert spec.data_dir.exists()
-    assert spec.bundled_packages_dir.exists()
-    assert (spec.data_dir / "data.txt").exists()
-    assert spec.requirements == requirements
-    assert (spec.bundled_packages_dir / "dep_pkg" / "__init__.py").exists()
-    assert (spec.bundled_packages_dir / "dep_pkg" / "file.py").exists()
-
-
 def test_truss_init_with_python_dx(tmp_path):
-    dir_name = str(tmp_path)
-    init_directory(dir_name, model_name="Test Model Name", python_configuration=True)
+    init_directory(tmp_path, model_name="Test Model Name", python_config=True)
 
     generated_truss_dir = code_gen.gen_truss_model_from_source(tmp_path / "my_model.py")
     truss_handle = load(generated_truss_dir)

truss/truss_handle/build.py

@@ -2,7 +2,7 @@
 import os
 import sys
 from pathlib import Path
-from typing import List, Optional, Union
+from typing import Optional, Union
 
 import yaml
 
@@ -61,11 +61,21 @@ def _populate_python_dx_target_directory(config: TrussConfig, dir_path: Path):
 
 
 def init_directory(
-    target_directory: str,
+    target_directory: Union[str, Path],
     build_config: Optional[Build] = None,
     model_name: Optional[str] = None,
-    python_configuration: bool = False,
+    python_config: bool = False,
 ) -> Path:
+    """
+    Initialize an empty placeholder Truss. A Truss is a build context designed
+    to be built as a container locally or uploaded into a baseten serving
+    environment. This placeholder structure can be filled to represent ML
+    models.
+
+    Args:
+        target_directory: Absolute or relative path of the directory to create
+                          Truss in. The directory is created if it doesn't exist.
+    """
     config = TrussConfig(
         model_name=model_name, python_version=map_local_to_supported_python_version()
     )
@@ -76,7 +86,7 @@ def init_directory(
     target_directory_path = Path(target_directory)
     target_directory_path.mkdir(parents=True, exist_ok=True)
 
-    if not python_configuration:
+    if not python_config:
         _populate_traditional_target_directory(
             config=config, dir_path=target_directory_path
         )
@@ -88,35 +98,6 @@ def init_directory(
     return target_directory_path
 
 
-def init(
-    target_directory: str,
-    data_files: Optional[List[str]] = None,
-    requirements_file: Optional[str] = None,
-    bundled_packages: Optional[List[str]] = None,
-    build_config: Optional[Build] = None,
-    model_name: Optional[str] = None,
-) -> TrussHandle:
-    """
-    Initialize an empty placeholder Truss. A Truss is a build context designed
-    to be built as a container locally or uploaded into a baseten serving
-    environment. This placeholder structure can be filled to represent ML
-    models.
-
-    Args:
-        target_directory: Absolute or relative path of the directory to create
-                          Truss in. The directory is created if it doesn't exist.
-    """
-    target_path = init_directory(
-        target_directory=target_directory,
-        build_config=build_config,
-        model_name=model_name,
-    )
-
-    th = TrussHandle(target_path)
-    _update_truss_props(th, data_files, requirements_file, bundled_packages)
-    return th
-
-
 def load(truss_directory: Union[str, Path]) -> TrussHandle:
     """Get a handle to a Truss. A Truss is a build context designed to be built
     as a container locally or uploaded into a model serving environment.
@@ -139,21 +120,3 @@ def cleanup() -> None:
             if (not obj.name == "config.yaml") and (obj.is_file()):
                 os.remove(obj)
     return
-
-
-def _update_truss_props(
-    scaf: TrussHandle,
-    data_files: Optional[List[str]] = None,
-    requirements_file: Optional[str] = None,
-    bundled_packages: Optional[List[str]] = None,
-) -> None:
-    if data_files is not None:
-        for data_file in data_files:
-            scaf.add_data(data_file)
-
-    if bundled_packages is not None:
-        for package in bundled_packages:
-            scaf.add_bundled_package(package)
-
-    if requirements_file is not None:
-        scaf.update_requirements_from_file(requirements_file)