snowflake-ml-python 1.4.1__py3-none-any.whl → 1.5.1__py3-none-any.whl

snowflake/ml/model/_client/sql/tag.py

@@ -1,52 +1,25 @@
 from typing import Any, Dict, List, Optional
 
-from snowflake.ml._internal.utils import (
-    identifier,
-    query_result_checker,
-    sql_identifier,
-)
-from snowflake.snowpark import row, session
+from snowflake.ml._internal.utils import query_result_checker, sql_identifier
+from snowflake.ml.model._client.sql import _base
+from snowflake.snowpark import row
 
 
-class ModuleTagSQLClient:
-    def __init__(
-        self,
-        session: session.Session,
-        *,
-        database_name: sql_identifier.SqlIdentifier,
-        schema_name: sql_identifier.SqlIdentifier,
-    ) -> None:
-        self._session = session
-        self._database_name = database_name
-        self._schema_name = schema_name
-
-    def __eq__(self, __value: object) -> bool:
-        if not isinstance(__value, ModuleTagSQLClient):
-            return False
-        return self._database_name == __value._database_name and self._schema_name == __value._schema_name
-
-    def fully_qualified_module_name(
-        self,
-        module_name: sql_identifier.SqlIdentifier,
-    ) -> str:
-        return identifier.get_schema_level_object_identifier(
-            self._database_name.identifier(), self._schema_name.identifier(), module_name.identifier()
-        )
-
+class ModuleTagSQLClient(_base._BaseSQLClient):
     def set_tag_on_model(
         self,
-        model_name: sql_identifier.SqlIdentifier,
         *,
-        tag_database_name: sql_identifier.SqlIdentifier,
-        tag_schema_name: sql_identifier.SqlIdentifier,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        tag_database_name: Optional[sql_identifier.SqlIdentifier],
+        tag_schema_name: Optional[sql_identifier.SqlIdentifier],
         tag_name: sql_identifier.SqlIdentifier,
         tag_value: str,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
-        fq_model_name = self.fully_qualified_module_name(model_name)
-        fq_tag_name = identifier.get_schema_level_object_identifier(
-            tag_database_name.identifier(), tag_schema_name.identifier(), tag_name.identifier()
-        )
+        fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
+        fq_tag_name = self.fully_qualified_object_name(tag_database_name, tag_schema_name, tag_name)
         query_result_checker.SqlResultValidator(
             self._session,
             f"ALTER MODEL {fq_model_name} SET TAG {fq_tag_name} = $${tag_value}$$",
@@ -55,17 +28,17 @@ class ModuleTagSQLClient:
 
     def unset_tag_on_model(
         self,
-        model_name: sql_identifier.SqlIdentifier,
         *,
-        tag_database_name: sql_identifier.SqlIdentifier,
-        tag_schema_name: sql_identifier.SqlIdentifier,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        tag_database_name: Optional[sql_identifier.SqlIdentifier],
+        tag_schema_name: Optional[sql_identifier.SqlIdentifier],
         tag_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
-        fq_model_name = self.fully_qualified_module_name(model_name)
-        fq_tag_name = identifier.get_schema_level_object_identifier(
-            tag_database_name.identifier(), tag_schema_name.identifier(), tag_name.identifier()
-        )
+        fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
+        fq_tag_name = self.fully_qualified_object_name(tag_database_name, tag_schema_name, tag_name)
         query_result_checker.SqlResultValidator(
             self._session,
             f"ALTER MODEL {fq_model_name} UNSET TAG {fq_tag_name}",
@@ -74,21 +47,21 @@ class ModuleTagSQLClient:
 
     def get_tag_value(
         self,
-        module_name: sql_identifier.SqlIdentifier,
         *,
-        tag_database_name: sql_identifier.SqlIdentifier,
-        tag_schema_name: sql_identifier.SqlIdentifier,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        tag_database_name: Optional[sql_identifier.SqlIdentifier],
+        tag_schema_name: Optional[sql_identifier.SqlIdentifier],
         tag_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> row.Row:
-        fq_module_name = self.fully_qualified_module_name(module_name)
-        fq_tag_name = identifier.get_schema_level_object_identifier(
-            tag_database_name.identifier(), tag_schema_name.identifier(), tag_name.identifier()
-        )
+        fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
+        fq_tag_name = self.fully_qualified_object_name(tag_database_name, tag_schema_name, tag_name)
         return (
             query_result_checker.SqlResultValidator(
                 self._session,
-                f"SELECT SYSTEM$GET_TAG($${fq_tag_name}$$, $${fq_module_name}$$, 'MODULE') AS TAG_VALUE",
+                f"SELECT SYSTEM$GET_TAG($${fq_tag_name}$$, $${fq_model_name}$$, 'MODULE') AS TAG_VALUE",
                 statement_params=statement_params,
             )
             .has_dimensions(expected_rows=1, expected_cols=1)
@@ -98,16 +71,19 @@ class ModuleTagSQLClient:
 
     def get_tag_list(
         self,
-        module_name: sql_identifier.SqlIdentifier,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> List[row.Row]:
-        fq_module_name = self.fully_qualified_module_name(module_name)
+        fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
+        actual_database_name = database_name or self._database_name
         return (
             query_result_checker.SqlResultValidator(
                 self._session,
                 f"""SELECT TAG_DATABASE, TAG_SCHEMA, TAG_NAME, TAG_VALUE
-FROM TABLE({self._database_name.identifier()}.INFORMATION_SCHEMA.TAG_REFERENCES($${fq_module_name}$$, 'MODULE'))""",
+FROM TABLE({actual_database_name.identifier()}.INFORMATION_SCHEMA.TAG_REFERENCES($${fq_model_name}$$, 'MODULE'))""",
                 statement_params=statement_params,
             )
             .has_column("TAG_DATABASE", allow_empty=True)

snowflake/ml/model/_deploy_client/image_builds/server_image_builder.py

@@ -37,6 +37,7 @@ class ServerImageBuilder(base_image_builder.ImageBuilder):
         session: snowpark.Session,
         artifact_stage_location: str,
         compute_pool: str,
+        job_name: str,
         external_access_integrations: List[str],
     ) -> None:
         """Initialization
@@ -49,6 +50,7 @@ class ServerImageBuilder(base_image_builder.ImageBuilder):
             artifact_stage_location: Spec file and future deployment related artifacts will be stored under
                 {stage}/models/{model_id}
             compute_pool: The compute pool used to run docker image build workload.
+            job_name: job_name to use.
             external_access_integrations: EAIs for network connection.
         """
         self.context_dir = context_dir
@@ -58,6 +60,7 @@ class ServerImageBuilder(base_image_builder.ImageBuilder):
         self.artifact_stage_location = artifact_stage_location
         self.compute_pool = compute_pool
         self.external_access_integrations = external_access_integrations
+        self.job_name = job_name
         self.client = snowservice_client.SnowServiceClient(session)
 
         assert artifact_stage_location.startswith(
@@ -203,8 +206,9 @@ class ServerImageBuilder(base_image_builder.ImageBuilder):
         )
 
     def _launch_kaniko_job(self, spec_stage_location: str) -> None:
-        logger.debug("Submitting job for building docker image with kaniko")
+        logger.debug(f"Submitting job {self.job_name} for building docker image with kaniko")
         self.client.create_job(
+            job_name=self.job_name,
             compute_pool=self.compute_pool,
             spec_stage_location=spec_stage_location,
             external_access_integrations=self.external_access_integrations,

snowflake/ml/model/_deploy_client/image_builds/templates/dockerfile_template

@@ -30,6 +30,7 @@ USER mambauser
 
 # Set MAMBA_DOCKERFILE_ACTIVATE=1 to activate the conda environment during build time.
 ARG MAMBA_DOCKERFILE_ACTIVATE=1
+ARG MAMBA_NO_LOW_SPEED_LIMIT=1
 
 # Bitsandbytes uses this ENVVAR to determine CUDA library location
 ENV CONDA_PREFIX=/opt/conda

snowflake/ml/model/_deploy_client/snowservice/deploy.py

@@ -346,6 +346,7 @@ class SnowServiceDeployment:
         (db, schema, _, _) = identifier.parse_schema_level_object_identifier(service_func_name)
 
         self._service_name = identifier.get_schema_level_object_identifier(db, schema, f"service_{model_id}")
+        self._job_name = identifier.get_schema_level_object_identifier(db, schema, f"build_{model_id}")
         # Spec file and future deployment related artifacts will be stored under {stage}/models/{model_id}
         self._model_artifact_stage_location = posixpath.join(deployment_stage_path, "models", self.id)
         self.debug_dir: Optional[str] = None
@@ -468,6 +469,7 @@ class SnowServiceDeployment:
                 session=self.session,
                 artifact_stage_location=self._model_artifact_stage_location,
                 compute_pool=self.options.compute_pool,
+                job_name=self._job_name,
                 external_access_integrations=self.options.external_access_integrations,
             )
         else:

snowflake/ml/model/_deploy_client/utils/constants.py

@@ -17,11 +17,6 @@ class ResourceStatus(Enum):
     INTERNAL_ERROR = "INTERNAL_ERROR"  # there was an internal service error.
 
 
-RESOURCE_TO_STATUS_FUNCTION_MAPPING = {
-    ResourceType.SERVICE: "SYSTEM$GET_SERVICE_STATUS",
-    ResourceType.JOB: "SYSTEM$GET_JOB_STATUS",
-}
-
 PREDICT = "predict"
 STAGE = "stage"
 COMPUTE_POOL = "compute_pool"

snowflake/ml/model/_deploy_client/utils/snowservice_client.py

@@ -70,13 +70,16 @@ class SnowServiceClient:
         logger.debug(f"Create service with SQL: \n {sql}")
         self.session.sql(sql).collect()
 
-    def create_job(self, compute_pool: str, spec_stage_location: str, external_access_integrations: List[str]) -> None:
+    def create_job(
+        self, job_name: str, compute_pool: str, spec_stage_location: str, external_access_integrations: List[str]
+    ) -> None:
         """Execute the job creation SQL command. Note that the job creation is synchronous, hence we execute it in a
         async way so that we can query the log in the meantime.
 
         Upon job failure, full job container log will be logged.
 
         Args:
+            job_name: name of the job
             compute_pool: name of the compute pool
             spec_stage_location: path to the stage location where the spec is located at.
             external_access_integrations: EAIs for network connection.
@@ -84,19 +87,18 @@ class SnowServiceClient:
         stage, path = uri.get_stage_and_path(spec_stage_location)
         sql = textwrap.dedent(
             f"""
-            EXECUTE SERVICE
+            EXECUTE JOB SERVICE
             IN COMPUTE POOL {compute_pool}
             FROM {stage}
-            SPEC = '{path}'
+            SPECIFICATION_FILE = '{path}'
+            NAME = {job_name}
             EXTERNAL_ACCESS_INTEGRATIONS = ({', '.join(external_access_integrations)})
             """
         )
         logger.debug(f"Create job with SQL: \n {sql}")
-        cur = self.session._conn._conn.cursor()
-        cur.execute_async(sql)
-        job_id = cur._sfqid
+        self.session.sql(sql).collect_nowait()
         self.block_until_resource_is_ready(
-            resource_name=str(job_id),
+            resource_name=job_name,
             resource_type=constants.ResourceType.JOB,
             container_name=constants.KANIKO_CONTAINER_NAME,
             max_retries=240,
@@ -182,10 +184,7 @@ class SnowServiceClient:
         """
         assert resource_type == constants.ResourceType.SERVICE or resource_type == constants.ResourceType.JOB
         query_command = ""
-        if resource_type == constants.ResourceType.SERVICE:
-            query_command = f"CALL SYSTEM$GET_SERVICE_LOGS('{resource_name}', '0', '{container_name}')"
-        elif resource_type == constants.ResourceType.JOB:
-            query_command = f"CALL SYSTEM$GET_JOB_LOGS('{resource_name}', '{container_name}')"
+        query_command = f"CALL SYSTEM$GET_SERVICE_LOGS('{resource_name}', '0', '{container_name}')"
         logger.warning(
             f"Best-effort log streaming from SPCS will be enabled when python logging level is set to INFO."
             f"Alternatively, you can also query the logs by running the query '{query_command}'"
@@ -201,7 +200,7 @@ class SnowServiceClient:
                 )
                 lsp.process_new_logs(resource_log, log_level=logging.INFO)
 
-            status = self.get_resource_status(resource_name=resource_name, resource_type=resource_type)
+            status = self.get_resource_status(resource_name=resource_name)
 
             if resource_type == constants.ResourceType.JOB and status == constants.ResourceStatus.DONE:
                 return
@@ -246,52 +245,24 @@ class SnowServiceClient:
     def get_resource_log(
         self, resource_name: str, resource_type: constants.ResourceType, container_name: str
     ) -> Optional[str]:
-        if resource_type == constants.ResourceType.SERVICE:
-            try:
-                row = self.session.sql(
-                    f"CALL SYSTEM$GET_SERVICE_LOGS('{resource_name}', '0', '{container_name}')"
-                ).collect()
-                return str(row[0]["SYSTEM$GET_SERVICE_LOGS"])
-            except Exception:
-                return None
-        elif resource_type == constants.ResourceType.JOB:
-            try:
-                row = self.session.sql(f"CALL SYSTEM$GET_JOB_LOGS('{resource_name}', '{container_name}')").collect()
-                return str(row[0]["SYSTEM$GET_JOB_LOGS"])
-            except Exception:
-                return None
-        else:
-            raise snowml_exceptions.SnowflakeMLException(
-                error_code=error_codes.NOT_IMPLEMENTED,
-                original_exception=NotImplementedError(
-                    f"{resource_type.name} is not yet supported in get_resource_log function"
-                ),
-            )
-
-    def get_resource_status(
-        self, resource_name: str, resource_type: constants.ResourceType
-    ) -> Optional[constants.ResourceStatus]:
+        try:
+            row = self.session.sql(
+                f"CALL SYSTEM$GET_SERVICE_LOGS('{resource_name}', '0', '{container_name}')"
+            ).collect()
+            return str(row[0]["SYSTEM$GET_SERVICE_LOGS"])
+        except Exception:
+            return None
+
+    def get_resource_status(self, resource_name: str) -> Optional[constants.ResourceStatus]:
         """Get resource status.
 
         Args:
             resource_name: Name of the resource.
-            resource_type: Type of the resource.
-
-        Raises:
-            SnowflakeMLException: If resource type does not have a corresponding system function for querying status.
-            SnowflakeMLException: If corresponding status call failed.
 
         Returns:
             Optional[constants.ResourceStatus]: The status of the resource, or None if the resource status is empty.
         """
-        if resource_type not in constants.RESOURCE_TO_STATUS_FUNCTION_MAPPING:
-            raise snowml_exceptions.SnowflakeMLException(
-                error_code=error_codes.INVALID_ARGUMENT,
-                original_exception=ValueError(
-                    f"Status querying is not supported for resources of type '{resource_type}'."
-                ),
-            )
-        status_func = constants.RESOURCE_TO_STATUS_FUNCTION_MAPPING[resource_type]
+        status_func = "SYSTEM$GET_SERVICE_STATUS"
         try:
             row = self.session.sql(f"CALL {status_func}('{resource_name}');").collect()
         except Exception:

snowflake/ml/model/_model_composer/model_composer.py

@@ -8,8 +8,10 @@ from typing import Any, Dict, List, Optional
 
 from absl import logging
 from packaging import requirements
+from typing_extensions import deprecated
 
 from snowflake.ml._internal import env as snowml_env, env_utils, file_utils
+from snowflake.ml._internal.lineage import data_source, lineage_utils
 from snowflake.ml.model import model_signature, type_hints as model_types
 from snowflake.ml.model._model_composer.model_manifest import model_manifest
 from snowflake.ml.model._packager import model_packager
@@ -134,6 +136,7 @@ class ModelComposer:
             model_meta=self.packager.meta,
             model_file_rel_path=pathlib.PurePosixPath(self.model_file_rel_path),
             options=options,
+            data_sources=self._get_data_sources(model),
         )
 
         file_utils.upload_directory_to_stage(
@@ -143,7 +146,8 @@ class ModelComposer:
             statement_params=self._statement_params,
         )
 
-    def load(
+    @deprecated("Only used by PrPr model registry. Use static method version of load instead.")
+    def legacy_load(
         self,
         *,
         meta_only: bool = False,
@@ -163,3 +167,20 @@ class ModelComposer:
         with zipfile.ZipFile(self.model_local_path, mode="r", compression=zipfile.ZIP_DEFLATED) as zf:
             zf.extractall(path=self._packager_workspace_path)
         self.packager.load(meta_only=meta_only, options=options)
+
+    @staticmethod
+    def load(
+        workspace_path: pathlib.Path,
+        *,
+        meta_only: bool = False,
+        options: Optional[model_types.ModelLoadOption] = None,
+    ) -> model_packager.ModelPackager:
+        mp = model_packager.ModelPackager(str(workspace_path / ModelComposer.MODEL_DIR_REL_PATH))
+        mp.load(meta_only=meta_only, options=options)
+        return mp
+
+    def _get_data_sources(self, model: model_types.SupportedModelType) -> Optional[List[data_source.DataSource]]:
+        data_sources = getattr(model, lineage_utils.DATA_SOURCES_ATTR, None)
+        if isinstance(data_sources, list) and all(isinstance(item, data_source.DataSource) for item in data_sources):
+            return data_sources
+        return None

snowflake/ml/model/_model_composer/model_manifest/model_manifest.py

@@ -5,6 +5,7 @@ from typing import List, Optional, cast
 
 import yaml
 
+from snowflake.ml._internal.lineage import data_source
 from snowflake.ml.model import type_hints
 from snowflake.ml.model._model_composer.model_manifest import model_manifest_schema
 from snowflake.ml.model._model_composer.model_method import (
@@ -36,6 +37,7 @@ class ModelManifest:
         model_meta: model_meta_api.ModelMetadata,
         model_file_rel_path: pathlib.PurePosixPath,
         options: Optional[type_hints.ModelSaveOption] = None,
+        data_sources: Optional[List[data_source.DataSource]] = None,
     ) -> None:
         if options is None:
             options = {}
@@ -90,6 +92,10 @@ class ModelManifest:
             ],
         )
 
+        lineage_sources = self._extract_lineage_info(data_sources)
+        if lineage_sources:
+            manifest_dict["lineage_sources"] = lineage_sources
+
         with (self.workspace_path / ModelManifest.MANIFEST_FILE_REL_PATH).open("w", encoding="utf-8") as f:
             # Anchors are not supported in the server, avoid that.
             yaml.SafeDumper.ignore_aliases = lambda *args: True  # type: ignore[method-assign]
@@ -108,3 +114,19 @@ class ModelManifest:
         res = cast(model_manifest_schema.ModelManifestDict, raw_input)
 
         return res
+
+    def _extract_lineage_info(
+        self, data_sources: Optional[List[data_source.DataSource]]
+    ) -> List[model_manifest_schema.LineageSourceDict]:
+        result = []
+        if data_sources:
+            for source in data_sources:
+                result.append(
+                    model_manifest_schema.LineageSourceDict(
+                        # Currently, we only support lineage from Dataset.
+                        type=model_manifest_schema.LineageSourceTypes.DATASET.value,
+                        entity=source.fully_qualified_name,
+                        version=source.version,
+                    )
+                )
+        return result

snowflake/ml/model/_model_composer/model_manifest/model_manifest_schema.py

@@ -75,8 +75,19 @@ class SnowparkMLDataDict(TypedDict):
     functions: Required[List[ModelFunctionInfoDict]]
 
 
+class LineageSourceTypes(enum.Enum):
+    DATASET = "DATASET"
+
+
+class LineageSourceDict(TypedDict):
+    type: Required[str]
+    entity: Required[str]
+    version: NotRequired[str]
+
+
 class ModelManifestDict(TypedDict):
     manifest_version: Required[str]
     runtimes: Required[Dict[str, ModelRuntimeDict]]
     methods: Required[List[ModelMethodDict]]
     user_data: NotRequired[Dict[str, Any]]
+    lineage_sources: NotRequired[List[LineageSourceDict]]

snowflake/ml/model/_packager/model_env/model_env.py

@@ -284,6 +284,7 @@ class ModelEnv:
                         " This may prevent model deploying to Snowflake Warehouse."
                     ),
                     category=UserWarning,
+                    stacklevel=2,
                 )
             if len(channel_dependencies) == 0 and channel not in self._conda_dependencies:
                 warnings.warn(
@@ -292,6 +293,7 @@ class ModelEnv:
                         " This may prevent model deploying to Snowflake Warehouse."
                     ),
                     category=UserWarning,
+                    stacklevel=2,
                 )
                 self._conda_dependencies[channel] = []
 
@@ -307,6 +309,7 @@ class ModelEnv:
                             " This may be unintentional."
                         ),
                         category=UserWarning,
+                        stacklevel=2,
                     )
 
         if pip_requirements_list:
@@ -316,6 +319,7 @@ class ModelEnv:
                     " This may prevent model deploying to Snowflake Warehouse."
                 ),
                 category=UserWarning,
+                stacklevel=2,
             )
             for pip_dependency in pip_requirements_list:
                 if any(
@@ -338,6 +342,7 @@ class ModelEnv:
                     " This may prevent model deploying to Snowflake Warehouse."
                 ),
                 category=UserWarning,
+                stacklevel=2,
             )
             for pip_dependency in pip_requirements_list:
                 if any(
@@ -372,3 +377,39 @@ class ModelEnv:
             "cuda_version": self.cuda_version,
             "snowpark_ml_version": self.snowpark_ml_version,
         }
+
+    def validate_with_local_env(
+        self, check_snowpark_ml_version: bool = False
+    ) -> List[env_utils.IncorrectLocalEnvironmentError]:
+        errors = []
+        try:
+            env_utils.validate_py_runtime_version(str(self._python_version))
+        except env_utils.IncorrectLocalEnvironmentError as e:
+            errors.append(e)
+
+        for conda_reqs in self._conda_dependencies.values():
+            for conda_req in conda_reqs:
+                try:
+                    env_utils.validate_local_installed_version_of_pip_package(
+                        env_utils.try_convert_conda_requirement_to_pip(conda_req)
+                    )
+                except env_utils.IncorrectLocalEnvironmentError as e:
+                    errors.append(e)
+
+        for pip_req in self._pip_requirements:
+            try:
+                env_utils.validate_local_installed_version_of_pip_package(pip_req)
+            except env_utils.IncorrectLocalEnvironmentError as e:
+                errors.append(e)
+
+        if check_snowpark_ml_version:
+            # For Modeling model
+            if self._snowpark_ml_version.base_version != snowml_env.VERSION:
+                errors.append(
+                    env_utils.IncorrectLocalEnvironmentError(
+                        f"The local installed version of Snowpark ML library is {snowml_env.VERSION} "
+                        f"which differs from required version {self.snowpark_ml_version}."
+                    )
+                )
+
+        return errors

snowflake/ml/model/_packager/model_handlers/mlflow.py

@@ -1,4 +1,5 @@
 import os
+import pathlib
 import tempfile
 from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
 
@@ -45,7 +46,7 @@ def _parse_mlflow_env(model_uri: str, env: model_env.ModelEnv) -> model_env.Mode
     if not os.path.exists(conda_env_file_path):
         raise ValueError("Cannot load MLFlow model dependencies.")
 
-    env.load_from_conda_file(conda_env_file_path)
+    env.load_from_conda_file(pathlib.Path(conda_env_file_path))
 
     return env
 

snowflake/ml/model/_packager/model_meta/model_meta.py

@@ -320,11 +320,7 @@ class ModelMetadata:
 
         with open(model_yaml_path, "w", encoding="utf-8") as out:
             yaml.SafeDumper.ignore_aliases = lambda *args: True  # type: ignore[method-assign]
-            yaml.safe_dump(
-                model_dict,
-                stream=out,
-                default_flow_style=False,
-            )
+            yaml.safe_dump(model_dict, stream=out, default_flow_style=False)
 
     @staticmethod
     def _validate_model_metadata(loaded_meta: Any) -> model_meta_schema.ModelMetadataDict:

snowflake/ml/model/_packager/model_packager.py

@@ -4,7 +4,6 @@ from typing import Dict, List, Optional
 
 from absl import logging
 
-from snowflake.ml._internal import env_utils
 from snowflake.ml._internal.exceptions import (
     error_codes,
     exceptions as snowml_exceptions,
@@ -129,8 +128,6 @@ class ModelPackager:
 
         model_meta.load_code_path(self.local_dir_path)
 
-        env_utils.validate_py_runtime_version(self.meta.env.python_version)
-
         handler = model_handler.load_handler(self.meta.model_type)
         if handler is None:
             raise snowml_exceptions.SnowflakeMLException(

snowflake/ml/modeling/_internal/local_implementations/pandas_trainer.py

@@ -3,6 +3,8 @@ from typing import List, Optional, Tuple
 
 import pandas as pd
 
+from snowflake.ml.modeling._internal.estimator_utils import handle_inference_result
+
 
 class PandasModelTrainer:
     """
@@ -72,11 +74,61 @@ class PandasModelTrainer:
             Tuple[pd.DataFrame, object]: [predicted dataset, estimator]
         """
         assert hasattr(self.estimator, "fit_predict")  # make type checker happy
-        args = {"X": self.dataset[self.input_cols]}
-        result = self.estimator.fit_predict(**args)
+        result = self.estimator.fit_predict(X=self.dataset[self.input_cols])
         result_df = pd.DataFrame(data=result, columns=expected_output_cols_list)
         if drop_input_cols:
             result_df = result_df
         else:
-            result_df = pd.concat([self.dataset, result_df], axis=1)
+            # in case the output column name overlap with the input column names,
+            # remove the ones in input column names
+            remove_dataset_col_name_exist_in_output_col = list(
+                set(self.dataset.columns) - set(expected_output_cols_list)
+            )
+            result_df = pd.concat([self.dataset[remove_dataset_col_name_exist_in_output_col], result_df], axis=1)
+        return (result_df, self.estimator)
+
+    def train_fit_transform(
+        self,
+        expected_output_cols_list: List[str],
+        drop_input_cols: Optional[bool] = False,
+    ) -> Tuple[pd.DataFrame, object]:
+        """Trains the model using specified features and target columns from the dataset.
+        This API is different from fit itself because it would also provide the transform
+        output.
+
+        Args:
+            expected_output_cols_list (List[str]): The output columns
+                name as a list. Defaults to None.
+            drop_input_cols (Optional[bool]): Boolean to determine whether to
+                drop the input columns from the output dataset.
+
+        Returns:
+            Tuple[pd.DataFrame, object]: [transformed dataset, estimator]
+        """
+        assert hasattr(self.estimator, "fit")  # make type checker happy
+        assert hasattr(self.estimator, "fit_transform")  # make type checker happy
+
+        argspec = inspect.getfullargspec(self.estimator.fit)
+        args = {"X": self.dataset[self.input_cols]}
+        if self.label_cols:
+            label_arg_name = "Y" if "Y" in argspec.args else "y"
+            args[label_arg_name] = self.dataset[self.label_cols].squeeze()
+
+        if self.sample_weight_col is not None and "sample_weight" in argspec.args:
+            args["sample_weight"] = self.dataset[self.sample_weight_col].squeeze()
+
+        inference_res = self.estimator.fit_transform(**args)
+
+        transformed_numpy_array, output_cols = handle_inference_result(
+            inference_res=inference_res, output_cols=expected_output_cols_list, inference_method="fit_transform"
+        )
+
+        result_df = pd.DataFrame(data=transformed_numpy_array, columns=output_cols)
+        if drop_input_cols:
+            result_df = result_df
+        else:
+            # in case the output column name overlap with the input column names,
+            # remove the ones in input column names
+            remove_dataset_col_name_exist_in_output_col = list(set(self.dataset.columns) - set(output_cols))
+            result_df = pd.concat([self.dataset[remove_dataset_col_name_exist_in_output_col], result_df], axis=1)
         return (result_df, self.estimator)