Map python versions greater than 3.9 to 3.9 (#128)

Author: pankajroark

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.1.6rc2"
+version = "0.1.6"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"

truss/build.py

@@ -10,7 +10,7 @@
 from truss.environment_inference.requirements_inference import infer_deps
 from truss.errors import FrameworkNotSupportedError
 from truss.model_frameworks import MODEL_FRAMEWORKS_BY_TYPE, model_framework_from_model
-from truss.model_inference import infer_python_version
+from truss.model_inference import infer_python_version, map_to_supported_python_version
 from truss.truss_config import DEFAULT_EXAMPLES_FILENAME, TrussConfig
 from truss.truss_handle import TrussHandle
 from truss.types import ModelFrameworkType
@@ -131,7 +131,7 @@ def mk_truss_from_pipeline(
     config = TrussConfig(
         model_type="custom",
         model_framework=ModelFrameworkType.CUSTOM,
-        python_version=infer_python_version(),
+        python_version=map_to_supported_python_version(infer_python_version()),
         requirements=requirements,
     )
 
@@ -225,7 +225,7 @@ def init(
     config = TrussConfig(
         model_type="custom",
         model_framework=ModelFrameworkType.CUSTOM,
-        python_version=infer_python_version(),
+        python_version=map_to_supported_python_version(infer_python_version()),
     )
 
     target_directory_path = populate_target_directory(

truss/model_framework.py

@@ -5,7 +5,7 @@
 import yaml
 from truss.constants import CONFIG_FILE, TEMPLATES_DIR
 from truss.environment_inference.requirements_inference import infer_deps
-from truss.model_inference import infer_python_version
+from truss.model_inference import infer_python_version, map_to_supported_python_version
 from truss.truss_config import DEFAULT_EXAMPLES_FILENAME, TrussConfig
 from truss.types import ModelFrameworkType
 from truss.utils import copy_file_path, copy_tree_path
@@ -58,7 +58,7 @@ def to_truss(self, model, target_directory: Path) -> str:
         else:
             target_examples_path.touch()
 
-        python_version = infer_python_version()
+        python_version = map_to_supported_python_version(infer_python_version())
 
         # Create config
         config = TrussConfig(

truss/model_inference.py

@@ -1,4 +1,5 @@
 import inspect
+import logging
 import pathlib
 import sys
 from ast import ClassDef, FunctionDef
@@ -25,20 +26,7 @@
     "py39",
 }
 
-
-def _get_entries_for_packages(list_of_requirements, desired_requirements):
-    name_to_req_str = {}
-    for req_name in desired_requirements:
-        for req_spec_full_str in list_of_requirements:
-            if "==" in req_spec_full_str:
-                req_spec_name, req_version = req_spec_full_str.split("==")
-                req_version_base = req_version.split("+")[0]
-                if req_name == req_spec_name:
-                    name_to_req_str[req_name] = f"{req_name}=={req_version_base}"
-            else:
-                continue
-
-    return name_to_req_str
+logger = logging.getLogger(__name__)
 
 
 def _infer_model_framework(model_class: str):
@@ -82,11 +70,40 @@ def _model_class(model: Any):
 
 
 def infer_python_version() -> str:
-    python_major_minor = f"py{sys.version_info.major}{sys.version_info.minor}"
-    # might want to fix up this logic
-    if python_major_minor not in PYTHON_VERSIONS:
-        python_major_minor = None
-    return python_major_minor
+    return f"py{sys.version_info.major}{sys.version_info.minor}"
+
+
+def map_to_supported_python_version(python_version: str) -> str:
+    """Map python version to truss supported python version.
+
+    Currently, it maps any versions greater than 3.9 to 3.9.
+
+    Args:
+        python_version: in the form py[major_version][minor_version] e.g. py39,
+        py310
+    """
+    python_major_version = int(python_version[2:3])
+    python_minor_version = int(python_version[3:])
+
+    if python_major_version > 3:
+        raise NotImplementedError("Only python version 3 is supported")
+
+    # TODO(pankaj) Add full support for 3.10 and 3.11, this is stop-gap.
+    if python_minor_version > 9:
+        logger.info(
+            f"Mapping python version {python_major_version}.{python_minor_version}"
+            " to 3.9, the highest version that Truss currently supports."
+        )
+        return "py39"
+
+    if python_minor_version < 7:
+        logger.info(
+            f"Mapping python version {python_major_version}.{python_minor_version}"
+            " to 3.7, the lowest version that Truss currently supports."
+        )
+        return "py37"
+
+    return python_version
 
 
 def infer_model_information(model: Any) -> ModelBuildStageOne:

truss/tests/test_model_inference.py

@@ -2,6 +2,7 @@
 from truss.constants import PYTORCH
 from truss.model_inference import (
     infer_model_information,
+    map_to_supported_python_version,
     validate_provided_parameters_with_model,
 )
 
@@ -36,3 +37,19 @@ def test_infer_model_information(pytorch_model_with_init_args):
     model_info = infer_model_information(pytorch_model_with_init_args[0])
     assert model_info.model_framework == PYTORCH
     assert model_info.model_type == "MyModel"
+
+
+@pytest.mark.parametrize(
+    "python_version, expected_python_version",
+    [
+        ("py37", "py37"),
+        ("py38", "py38"),
+        ("py39", "py39"),
+        ("py310", "py39"),
+        ("py311", "py39"),
+        ("py36", "py37"),
+    ],
+)
+def test_map_to_supported_python_version(python_version, expected_python_version):
+    out_python_version = map_to_supported_python_version(python_version)
+    assert out_python_version == expected_python_version