snowflake-ml-python 1.0.1__py3-none-any.whl → 1.0.3__py3-none-any.whl

snowflake/ml/model/_model_meta.py

@@ -1,12 +1,13 @@
 import dataclasses
+import importlib
 import os
 import sys
 import warnings
+from collections import namedtuple
 from contextlib import contextmanager
 from datetime import datetime
-from pathlib import Path
 from types import ModuleType
-from typing import Any, Callable, Dict, Generator, List, Optional, Sequence, Tuple, cast
+from typing import Any, Callable, Dict, Generator, List, Optional, Sequence, cast
 
 import cloudpickle
 import yaml
@@ -24,7 +25,7 @@ from snowflake.snowpark import DataFrame as SnowparkDataFrame
 MODEL_METADATA_VERSION = 1
 _BASIC_DEPENDENCIES = _core_requirements.REQUIREMENTS
 
-_BASIC_DEPENDENCIES.append(env_utils._SNOWML_PKG_NAME)
+Dependency = namedtuple("Dependency", ["conda_name", "pip_name"])
 
 
 @dataclasses.dataclass
@@ -84,6 +85,10 @@ def _create_model_metadata(
         A model metadata object.
     """
     model_dir_path = os.path.normpath(model_dir_path)
+    embed_local_ml_library = kwargs.pop("embed_local_ml_library", False)
+    if embed_local_ml_library:
+        snowml_path = list(importlib.import_module("snowflake.ml").__path__)[0]
+        kwargs["local_ml_library_version"] = f"{snowml_env.VERSION}+{file_utils.hash_directory(snowml_path)}"
 
     model_meta = ModelMetadata(
         name=name,
@@ -100,6 +105,14 @@ def _create_model_metadata(
         os.makedirs(code_dir_path, exist_ok=True)
         for code_path in code_paths:
             file_utils.copy_file_or_tree(code_path, code_dir_path)
+
+    if embed_local_ml_library:
+        code_dir_path = os.path.join(model_dir_path, ModelMetadata.MODEL_CODE_DIR)
+        snowml_path = list(importlib.import_module("snowflake.ml").__path__)[0]
+        snowml_path_in_code = os.path.join(code_dir_path, "snowflake")
+        os.makedirs(snowml_path_in_code, exist_ok=True)
+        file_utils.copy_file_or_tree(snowml_path, snowml_path_in_code)
+
     try:
         imported_modules = []
         if ext_modules:
@@ -117,8 +130,7 @@ def _create_model_metadata(
 
 def _load_model_metadata(model_dir_path: str) -> "ModelMetadata":
     """Load models for a directory. Model is initially loaded normally. If additional codes are included when packed,
-        the code path is added to system path to be imported and overwrites those modules with the same name that has
-        been imported.
+        the code path is added to system path to be imported with highest priority.
 
     Args:
         model_dir_path: Path to the directory containing the model to be loaded.
@@ -131,14 +143,12 @@ def _load_model_metadata(model_dir_path: str) -> "ModelMetadata":
     meta = ModelMetadata.load_model_metadata(model_dir_path)
     code_path = os.path.join(model_dir_path, ModelMetadata.MODEL_CODE_DIR)
     if os.path.exists(code_path):
-        sys.path = [code_path] + sys.path
-        modules = [
-            p.stem
-            for p in Path(code_path).rglob("*.py")
-            if p.is_file() and p.name != "__init__.py" and p.name != "__main__.py"
-        ]
+        if code_path in sys.path:
+            sys.path.remove(code_path)
+        sys.path.insert(0, code_path)
+        modules = file_utils.get_all_modules(code_path)
         for module in modules:
-            sys.modules.pop(module, None)
+            sys.modules.pop(module.name, None)
     return meta
 
 
@@ -206,8 +216,12 @@ class ModelMetadata:
         self._pip_requirements = env_utils.validate_pip_requirement_string_list(
             pip_requirements if pip_requirements else []
         )
-
-        self._include_if_absent([(dep, dep) for dep in _BASIC_DEPENDENCIES])
+        if "local_ml_library_version" in kwargs:
+            self._include_if_absent([Dependency(conda_name=dep, pip_name=dep) for dep in _BASIC_DEPENDENCIES])
+        else:
+            self._include_if_absent(
+                [Dependency(conda_name=dep, pip_name=dep) for dep in _BASIC_DEPENDENCIES + [env_utils._SNOWML_PKG_NAME]]
+            )
 
         self.__dict__.update(kwargs)
 
@@ -225,7 +239,7 @@ class ModelMetadata:
             for req in reqs
         )
 
-    def _include_if_absent(self, pkgs: List[Tuple[str, str]]) -> None:
+    def _include_if_absent(self, pkgs: List[Dependency]) -> None:
         conda_reqs_str, pip_reqs_str = tuple(zip(*pkgs))
         pip_reqs = env_utils.validate_pip_requirement_string_list(list(pip_reqs_str))
         conda_reqs = env_utils.validate_conda_dependency_string_list(list(conda_reqs_str))
@@ -318,7 +332,7 @@ class ModelMetadata:
             path: The path of the directory to write a yaml file in it.
         """
         model_yaml_path = os.path.join(path, ModelMetadata.MODEL_METADATA_FILE)
-        with open(model_yaml_path, "w") as out:
+        with open(model_yaml_path, "w", encoding="utf-8") as out:
             yaml.safe_dump({**self.to_dict(), "version": MODEL_METADATA_VERSION}, stream=out, default_flow_style=False)
 
         env_dir_path = os.path.join(path, ModelMetadata.ENV_DIR)
@@ -341,10 +355,10 @@ class ModelMetadata:
             Loaded model metadata object.
         """
         model_yaml_path = os.path.join(path, ModelMetadata.MODEL_METADATA_FILE)
-        with open(model_yaml_path) as f:
+        with open(model_yaml_path, encoding="utf-8") as f:
             loaded_mata = yaml.safe_load(f.read())
 
-        loaded_mata_version = loaded_mata.get("version", None)
+        loaded_mata_version = loaded_mata.pop("version", None)
         if not loaded_mata_version or loaded_mata_version != MODEL_METADATA_VERSION:
             raise NotImplementedError("Unknown or unsupported model metadata file found.")
 
@@ -383,7 +397,7 @@ def _validate_signature(
     if isinstance(sample_input, SnowparkDataFrame):
         # Added because of Any from missing stubs.
         trunc_sample_input = cast(SnowparkDataFrame, trunc_sample_input)
-        local_sample_input = trunc_sample_input.to_pandas()
+        local_sample_input = model_signature._SnowparkDataFrameHandler.convert_to_df(trunc_sample_input)
     else:
         local_sample_input = trunc_sample_input
     for target_method in target_methods: