snowflake-ml-python 1.9.1__tar.gz → 1.9.2__tar.gz

{snowflake_ml_python-1.9.1 → snowflake_ml_python-1.9.2}/CHANGELOG.md

@@ -1,6 +1,43 @@
 # Release History
 
-## 1.9.1
+## 1.9.2
+
+### Bug Fixes
+
+- DataConnector: Fix `self._session` related errors inside Container Runtime.
+- Registry: Fix a bug when trying to pass `None` to array (`pd.dtype('O')`) in signature and pandas data handler.
+
+### New Features
+
+- Experiment Tracking (PrPr): Automatically log the model, metrics, and parameters while training
+  XGBoost and LightGBM models.
+
+```python
+from snowflake.ml.experiment import ExperimentTracking
+from snowflake.ml.experiment.callback import SnowflakeXgboostCallback, SnowflakeLightgbmCallback
+
+exp = ExperimentTracking(session=sp_session, database_name="ML", schema_name="PUBLIC")
+
+exp.set_experiment("MY_EXPERIMENT")
+
+# XGBoost
+callback = SnowflakeXgboostCallback(
+  exp, log_model=True, log_metrics=True, log_params=True, model_name="model_name", model_signature=sig
+)
+model = XGBClassifier(callbacks=[callback])
+with exp.start_run():
+  model.fit(X, y, eval_set=[(X_test, y_test)])
+
+# LightGBM
+callback = SnowflakeLightgbmCallback(
+  exp, log_model=True, log_metrics=True, log_params=True, model_name="model_name", model_signature=sig
+)
+model = LGBMClassifier()
+with exp.start_run():
+  model.fit(X, y, eval_set=[(X_test, y_test)], callbacks=[callback])
+```
+
+## 1.9.1 (07-18-2025)
 
 ### Bug Fixes
 

{snowflake_ml_python-1.9.1/snowflake_ml_python.egg-info → snowflake_ml_python-1.9.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowflake-ml-python
-Version: 1.9.1
+Version: 1.9.2
 Summary: The machine learning client library that is used for interacting with Snowflake to build machine learning solutions.
 Author-email: "Snowflake, Inc" <support@snowflake.com>
 License: 
@@ -243,6 +243,7 @@ Requires-Dist: importlib_resources<7,>=6.1.1
 Requires-Dist: numpy<3,>=1.23
 Requires-Dist: packaging<25,>=20.9
 Requires-Dist: pandas<3,>=2.1.4
+Requires-Dist: platformdirs<5
 Requires-Dist: pyarrow
 Requires-Dist: pydantic<3,>=2.8.2
 Requires-Dist: pyjwt<3,>=2.0.0
@@ -273,7 +274,7 @@ Requires-Dist: tensorflow<3,>=2.17.0; extra == "all"
 Requires-Dist: tokenizers<1,>=0.15.1; extra == "all"
 Requires-Dist: torch<3,>=2.0.1; extra == "all"
 Requires-Dist: torchdata<1,>=0.4; extra == "all"
-Requires-Dist: transformers<5,>=4.39.3; extra == "all"
+Requires-Dist: transformers!=4.51.3,<5,>=4.39.3; extra == "all"
 Provides-Extra: altair
 Requires-Dist: altair<6,>=5; extra == "altair"
 Provides-Extra: catboost
@@ -298,7 +299,7 @@ Requires-Dist: sentence-transformers<4,>=2.7.0; extra == "transformers"
 Requires-Dist: sentencepiece<0.2.0,>=0.1.95; extra == "transformers"
 Requires-Dist: tokenizers<1,>=0.15.1; extra == "transformers"
 Requires-Dist: torch<3,>=2.0.1; extra == "transformers"
-Requires-Dist: transformers<5,>=4.39.3; extra == "transformers"
+Requires-Dist: transformers!=4.51.3,<5,>=4.39.3; extra == "transformers"
 Dynamic: license-file
 
 # Snowpark ML
@@ -409,7 +410,44 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
 
 # Release History
 
-## 1.9.1
+## 1.9.2
+
+### Bug Fixes
+
+- DataConnector: Fix `self._session` related errors inside Container Runtime.
+- Registry: Fix a bug when trying to pass `None` to array (`pd.dtype('O')`) in signature and pandas data handler.
+
+### New Features
+
+- Experiment Tracking (PrPr): Automatically log the model, metrics, and parameters while training
+  XGBoost and LightGBM models.
+
+```python
+from snowflake.ml.experiment import ExperimentTracking
+from snowflake.ml.experiment.callback import SnowflakeXgboostCallback, SnowflakeLightgbmCallback
+
+exp = ExperimentTracking(session=sp_session, database_name="ML", schema_name="PUBLIC")
+
+exp.set_experiment("MY_EXPERIMENT")
+
+# XGBoost
+callback = SnowflakeXgboostCallback(
+  exp, log_model=True, log_metrics=True, log_params=True, model_name="model_name", model_signature=sig
+)
+model = XGBClassifier(callbacks=[callback])
+with exp.start_run():
+  model.fit(X, y, eval_set=[(X_test, y_test)])
+
+# LightGBM
+callback = SnowflakeLightgbmCallback(
+  exp, log_model=True, log_metrics=True, log_params=True, model_name="model_name", model_signature=sig
+)
+model = LGBMClassifier()
+with exp.start_run():
+  model.fit(X, y, eval_set=[(X_test, y_test)], callbacks=[callback])
+```
+
+## 1.9.1 (07-18-2025)
 
 ### Bug Fixes
 

{snowflake_ml_python-1.9.1 → snowflake_ml_python-1.9.2}/pyproject.toml

@@ -8,7 +8,7 @@ description = "The machine learning client library that is used for interacting
 classifiers = [ "Development Status :: 5 - Production/Stable", "Environment :: Console", "Environment :: Other Environment", "Intended Audience :: Developers", "Intended Audience :: Education", "Intended Audience :: Information Technology", "Intended Audience :: System Administrators", "License :: OSI Approved :: Apache Software License", "Operating System :: OS Independent", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Topic :: Database", "Topic :: Software Development", "Topic :: Software Development :: Libraries", "Topic :: Software Development :: Libraries :: Application Frameworks", "Topic :: Software Development :: Libraries :: Python Modules", "Topic :: Scientific/Engineering :: Information Analysis",]
 requires-python = ">=3.9, <3.13"
 dynamic = [ "version", "readme",]
-dependencies = [ "absl-py>=0.15,<2", "anyio>=3.5.0,<5", "cachetools>=3.1.1,<6", "cloudpickle>=2.0.0", "cryptography", "fsspec[http]>=2024.6.1,<2026", "importlib_resources>=6.1.1, <7", "numpy>=1.23,<3", "packaging>=20.9,<25", "pandas>=2.1.4,<3", "pyarrow", "pydantic>=2.8.2, <3", "pyjwt>=2.0.0, <3", "pytimeparse>=1.1.8,<2", "pyyaml>=6.0,<7", "retrying>=1.3.3,<2", "s3fs>=2024.6.1,<2026", "scikit-learn<1.6", "scipy>=1.9,<2", "shap>=0.46.0,<1", "snowflake-connector-python[pandas]>=3.15.0,<4", "snowflake-snowpark-python>=1.17.0,<2,!=1.26.0", "snowflake.core>=1.0.2,<2", "sqlparse>=0.4,<1", "tqdm<5", "typing-extensions>=4.1.0,<5", "xgboost>=1.7.3,<3",]
+dependencies = [ "absl-py>=0.15,<2", "anyio>=3.5.0,<5", "cachetools>=3.1.1,<6", "cloudpickle>=2.0.0", "cryptography", "fsspec[http]>=2024.6.1,<2026", "importlib_resources>=6.1.1, <7", "numpy>=1.23,<3", "packaging>=20.9,<25", "pandas>=2.1.4,<3", "platformdirs<5", "pyarrow", "pydantic>=2.8.2, <3", "pyjwt>=2.0.0, <3", "pytimeparse>=1.1.8,<2", "pyyaml>=6.0,<7", "retrying>=1.3.3,<2", "s3fs>=2024.6.1,<2026", "scikit-learn<1.6", "scipy>=1.9,<2", "shap>=0.46.0,<1", "snowflake-connector-python[pandas]>=3.15.0,<4", "snowflake-snowpark-python>=1.17.0,<2,!=1.26.0", "snowflake.core>=1.0.2,<2", "sqlparse>=0.4,<1", "tqdm<5", "typing-extensions>=4.1.0,<5", "xgboost>=1.7.3,<3",]
 [[project.authors]]
 name = "Snowflake, Inc"
 email = "support@snowflake.com"
@@ -24,7 +24,7 @@ Issues = "https://github.com/snowflakedb/snowflake-ml-python/issues"
 Changelog = "https://github.com/snowflakedb/snowflake-ml-python/blob/master/CHANGELOG.md"
 
 [project.optional-dependencies]
-all = [ "altair>=5,<6", "catboost>=1.2.0, <2", "keras>=2.0.0,<4", "lightgbm>=4.1.0, <5", "mlflow>=2.16.0, <3", "sentence-transformers>=2.7.0,<4", "sentencepiece>=0.1.95,<0.2.0", "streamlit>=1.30.0,<2", "tensorflow>=2.17.0,<3", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "torchdata>=0.4,<1", "transformers>=4.39.3,<5",]
+all = [ "altair>=5,<6", "catboost>=1.2.0, <2", "keras>=2.0.0,<4", "lightgbm>=4.1.0, <5", "mlflow>=2.16.0, <3", "sentence-transformers>=2.7.0,<4", "sentencepiece>=0.1.95,<0.2.0", "streamlit>=1.30.0,<2", "tensorflow>=2.17.0,<3", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "torchdata>=0.4,<1", "transformers>=4.39.3,<5,!=4.51.3",]
 altair = [ "altair>=5,<6",]
 catboost = [ "catboost>=1.2.0, <2",]
 keras = [ "keras>=2.0.0,<4", "tensorflow>=2.17.0,<3", "torch>=2.0.1,<3",]
@@ -33,7 +33,7 @@ mlflow = [ "mlflow>=2.16.0, <3",]
 streamlit = [ "streamlit>=1.30.0,<2",]
 tensorflow = [ "tensorflow>=2.17.0,<3",]
 torch = [ "torch>=2.0.1,<3", "torchdata>=0.4,<1",]
-transformers = [ "sentence-transformers>=2.7.0,<4", "sentencepiece>=0.1.95,<0.2.0", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "transformers>=4.39.3,<5",]
+transformers = [ "sentence-transformers>=2.7.0,<4", "sentencepiece>=0.1.95,<0.2.0", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "transformers>=4.39.3,<5,!=4.51.3",]
 
 [tool.setuptools.package-data]
 "*" = [ "*",]

{snowflake_ml_python-1.9.1 → snowflake_ml_python-1.9.2}/snowflake/ml/_internal/utils/mixins.py

@@ -21,10 +21,12 @@ class SerializableSessionMixin:
 
     def __getstate__(self) -> dict[str, Any]:
         """Customize pickling to exclude non-serializable session and related components."""
-        if hasattr(super(), "__getstate__"):
-            state: dict[str, Any] = super().__getstate__()  # type: ignore[misc]
-        else:
-            state = self.__dict__.copy()
+        parent_state = (
+            super().__getstate__()  # type: ignore[misc] # object.__getstate__ appears in 3.11
+            if hasattr(super(), "__getstate__")
+            else self.__dict__
+        )
+        state = dict(parent_state)  # Create a copy so we can safely modify the state
 
         # Save session metadata for validation during unpickling
         session = state.pop(_SESSION_KEY, None)

snowflake_ml_python-1.9.2/snowflake/ml/_internal/utils/service_logger.py

@@ -0,0 +1,165 @@
+import enum
+import logging
+import os
+import sys
+import tempfile
+import time
+import uuid
+from typing import Optional
+
+import platformdirs
+
+
+class LogColor(enum.Enum):
+    GREY = "\x1b[38;20m"
+    RED = "\x1b[31;20m"
+    BOLD_RED = "\x1b[31;1m"
+    YELLOW = "\x1b[33;20m"
+    BLUE = "\x1b[34;20m"
+    GREEN = "\x1b[32;20m"
+    ORANGE = "\x1b[38;5;214m"
+    BOLD_ORANGE = "\x1b[38;5;214;1m"
+    PURPLE = "\x1b[35;20m"
+    BOLD_PURPLE = "\x1b[35;1m"
+
+
+class CustomFormatter(logging.Formatter):
+
+    reset = "\x1b[0m"
+    log_format = "%(name)s [%(asctime)s] [%(levelname)s] %(message)s"
+
+    def __init__(self, info_color: LogColor) -> None:
+        super().__init__()
+        self.level_colors = {
+            logging.DEBUG: LogColor.GREY.value,
+            logging.INFO: info_color.value,
+            logging.WARNING: LogColor.YELLOW.value,
+            logging.ERROR: LogColor.RED.value,
+            logging.CRITICAL: LogColor.BOLD_RED.value,
+        }
+
+    def format(self, record: logging.LogRecord) -> str:
+        # default to DEBUG color
+        fmt = self.level_colors.get(record.levelno, self.level_colors[logging.DEBUG]) + self.log_format + self.reset
+        formatter = logging.Formatter(fmt)
+
+        # split the log message by lines and format each line individually
+        original_message = record.getMessage()
+        message_lines = original_message.splitlines()
+        formatted_lines = [
+            formatter.format(
+                logging.LogRecord(
+                    name=record.name,
+                    level=record.levelno,
+                    pathname=record.pathname,
+                    lineno=record.lineno,
+                    msg=line,
+                    args=None,
+                    exc_info=None,
+                )
+            )
+            for line in message_lines
+        ]
+
+        return "\n".join(formatted_lines)
+
+
+def _test_writability(directory: str) -> bool:
+    """Test if a directory is writable by creating and removing a test file."""
+    try:
+        os.makedirs(directory, exist_ok=True)
+        test_file = os.path.join(directory, f".write_test_{uuid.uuid4().hex[:8]}")
+        with open(test_file, "w") as f:
+            f.write("test")
+        os.remove(test_file)
+        return True
+    except OSError:
+        return False
+
+
+def _try_log_location(log_dir: str, operation_id: str) -> Optional[str]:
+    """Try to create a log file in the given directory if it's writable."""
+    if _test_writability(log_dir):
+        return os.path.join(log_dir, f"{operation_id}.log")
+    return None
+
+
+def _get_log_file_path(operation_id: str) -> Optional[str]:
+    """Get platform-independent log file path. Returns None if no writable location found."""
+    # Try locations in order of preference
+    locations = [
+        # Primary: User log directory
+        platformdirs.user_log_dir("snowflake-ml", "Snowflake"),
+        # Fallback 1: System temp directory
+        os.path.join(tempfile.gettempdir(), "snowflake-ml-logs"),
+        # Fallback 2: Current working directory
+        ".",
+    ]
+
+    for location in locations:
+        log_file_path = _try_log_location(location, operation_id)
+        if log_file_path:
+            return log_file_path
+
+    # No writable location found
+    return None
+
+
+def _get_or_create_parent_logger(operation_id: str) -> logging.Logger:
+    """Get or create a parent logger with FileHandler for the operation."""
+    parent_logger_name = f"snowflake_ml_operation_{operation_id}"
+    parent_logger = logging.getLogger(parent_logger_name)
+
+    # Only add handler if it doesn't exist yet
+    if not parent_logger.handlers:
+        log_file_path = _get_log_file_path(operation_id)
+
+        if log_file_path:
+            # Successfully found a writable location
+            try:
+                file_handler = logging.FileHandler(log_file_path)
+                file_handler.setFormatter(logging.Formatter("%(name)s [%(asctime)s] [%(levelname)s] %(message)s"))
+                parent_logger.addHandler(file_handler)
+                parent_logger.setLevel(logging.DEBUG)
+                parent_logger.propagate = False  # Don't propagate to root logger
+
+                # Log the file location
+                parent_logger.warning(f"Operation logs saved to: {log_file_path}")
+            except OSError as e:
+                # Even though we found a path, file creation failed
+                # Fall back to console-only logging
+                parent_logger.setLevel(logging.DEBUG)
+                parent_logger.propagate = False
+                parent_logger.warning(f"Could not create log file at {log_file_path}: {e}. Using console-only logging.")
+        else:
+            # No writable location found, use console-only logging
+            parent_logger.setLevel(logging.DEBUG)
+            parent_logger.propagate = False
+            parent_logger.warning("Filesystem appears to be readonly. Using console-only logging.")
+
+    return parent_logger
+
+
+def get_logger(logger_name: str, info_color: LogColor, operation_id: Optional[str] = None) -> logging.Logger:
+    logger = logging.getLogger(logger_name)
+    handler = logging.StreamHandler(sys.stdout)
+    handler.setFormatter(CustomFormatter(info_color))
+    logger.addHandler(handler)
+
+    # If operation_id provided, set up parent logger with file handler
+    if operation_id:
+        parent_logger = _get_or_create_parent_logger(operation_id)
+        logger.parent = parent_logger
+        logger.propagate = True
+
+    return logger
+
+
+def get_operation_id() -> str:
+    """Generate a unique operation ID."""
+    return f"model_deploy_{uuid.uuid4().hex[:8]}_{int(time.time())}"
+
+
+def get_log_file_location(operation_id: str) -> Optional[str]:
+    """Get the log file path for an operation ID. Returns None if no writable location available."""
+    return _get_log_file_path(operation_id)

{snowflake_ml_python-1.9.1 → snowflake_ml_python-1.9.2}/snowflake/ml/data/_internal/arrow_ingestor.py

@@ -14,6 +14,7 @@ if TYPE_CHECKING:
     import ray
 
 from snowflake import snowpark
+from snowflake.ml._internal.utils import mixins
 from snowflake.ml.data import data_ingestor, data_source, ingestor_utils
 
 _EMPTY_RECORD_BATCH = pa.RecordBatch.from_arrays([], [])
@@ -44,7 +45,7 @@ class _RecordBatchesBuffer:
         return popped
 
 
-class ArrowIngestor(data_ingestor.DataIngestor):
+class ArrowIngestor(data_ingestor.DataIngestor, mixins.SerializableSessionMixin):
     """Read and parse the data sources into an Arrow Dataset and yield batched numpy array in dict."""
 
     def __init__(
@@ -71,6 +72,8 @@ class ArrowIngestor(data_ingestor.DataIngestor):
 
     @classmethod
     def from_sources(cls, session: snowpark.Session, sources: Sequence[data_source.DataSource]) -> "ArrowIngestor":
+        if session is None:
+            raise ValueError("Session is required")
         return cls(session, sources)
 
     @classmethod

{snowflake_ml_python-1.9.1 → snowflake_ml_python-1.9.2}/snowflake/ml/data/data_connector.py

@@ -6,10 +6,9 @@ from typing_extensions import deprecated
 
 from snowflake import snowpark
 from snowflake.ml._internal import env, telemetry
-from snowflake.ml._internal.utils import mixins
 from snowflake.ml.data import data_ingestor, data_source
 from snowflake.ml.data._internal.arrow_ingestor import ArrowIngestor
-from snowflake.snowpark import context as sf_context
+from snowflake.snowpark import context as sp_context
 
 if TYPE_CHECKING:
     import pandas as pd
@@ -22,13 +21,11 @@ if TYPE_CHECKING:
     from snowflake.ml import dataset
 
 _PROJECT = "DataConnector"
-_INGESTOR_KEY = "_ingestor"
-_INGESTOR_SOURCES_KEY = "ingestor$sources"
 
 DataConnectorType = TypeVar("DataConnectorType", bound="DataConnector")
 
 
-class DataConnector(mixins.SerializableSessionMixin):
+class DataConnector:
     """Snowflake data reader which provides application integration connectors"""
 
     DEFAULT_INGESTOR_CLASS: type[data_ingestor.DataIngestor] = ArrowIngestor
@@ -36,11 +33,8 @@ class DataConnector(mixins.SerializableSessionMixin):
     def __init__(
         self,
         ingestor: data_ingestor.DataIngestor,
-        *,
-        session: Optional[snowpark.Session] = None,
         **kwargs: Any,
     ) -> None:
-        self._session = session
         self._ingestor = ingestor
         self._kwargs = kwargs
 
@@ -63,7 +57,7 @@ class DataConnector(mixins.SerializableSessionMixin):
         ingestor_class: Optional[type[data_ingestor.DataIngestor]] = None,
         **kwargs: Any,
     ) -> DataConnectorType:
-        session = session or sf_context.get_active_session()
+        session = session or sp_context.get_active_session()
         source = data_source.DataFrameInfo(query)
         return cls.from_sources(session, [source], ingestor_class=ingestor_class, **kwargs)
 
@@ -107,31 +101,7 @@ class DataConnector(mixins.SerializableSessionMixin):
     ) -> DataConnectorType:
         ingestor_class = ingestor_class or cls.DEFAULT_INGESTOR_CLASS
         ingestor = ingestor_class.from_sources(session, sources)
-        return cls(ingestor, **kwargs, session=session)
-
-    def __getstate__(self) -> dict[str, Any]:
-        """Customize pickling to exclude non-serializable session and related components."""
-        if hasattr(super(), "__getstate__"):
-            state = super().__getstate__()
-        else:
-            state = self.__dict__.copy()
-
-        ingestor = state.pop(_INGESTOR_KEY)
-        state[_INGESTOR_SOURCES_KEY] = ingestor.data_sources
-
-        return state
-
-    def __setstate__(self, state: dict[str, Any]) -> None:
-        """Restore session from context during unpickling."""
-        data_sources = state.pop(_INGESTOR_SOURCES_KEY)
-
-        if hasattr(super(), "__setstate__"):
-            super().__setstate__(state)
-        else:
-            self.__dict__.update(state)
-
-        assert self._session is not None
-        self._ingestor = self.DEFAULT_INGESTOR_CLASS.from_sources(self._session, data_sources)
+        return cls(ingestor, **kwargs)
 
     @property
     def data_sources(self) -> list[data_source.DataSource]:

{snowflake_ml_python-1.9.1 → snowflake_ml_python-1.9.2}/snowflake/ml/dataset/dataset.py

@@ -177,7 +177,7 @@ class Dataset(lineage_node.LineageNode):
                 original_exception=RuntimeError("No Dataset version selected."),
             )
         if self._reader is None:
-            self._reader = dataset_reader.DatasetReader.from_dataset(self)
+            self._reader = dataset_reader.DatasetReader.from_dataset(self, snowpark_session=self._session)
         return self._reader
 
     @staticmethod

{snowflake_ml_python-1.9.1 → snowflake_ml_python-1.9.2}/snowflake/ml/dataset/dataset_reader.py

@@ -1,5 +1,4 @@
 from typing import Any, Optional
-from warnings import warn
 
 from snowflake import snowpark
 from snowflake.ml._internal import telemetry
@@ -21,16 +20,11 @@ class DatasetReader(data_connector.DataConnector, mixins.SerializableSessionMixi
         self,
         ingestor: data_ingestor.DataIngestor,
         *,
-        session: snowpark.Session,
         snowpark_session: Optional[snowpark.Session] = None,
     ) -> None:
-        if snowpark_session is not None:
-            warn(
-                "Argument snowpark_session is deprecated and will be removed in a future release. Use session instead."
-            )
-            session = snowpark_session
-        super().__init__(ingestor, session=session)
+        super().__init__(ingestor)
 
+        self._session = snowpark_session
         self._fs_cached: Optional[snowfs.SnowFileSystem] = None
         self._files: Optional[list[str]] = None
 

snowflake_ml_python-1.9.2/snowflake/ml/experiment/__init__.py

@@ -0,0 +1,3 @@
+from snowflake.ml.experiment.experiment_tracking import ExperimentTracking
+
+__all__ = ["ExperimentTracking"]

snowflake_ml_python-1.9.2/snowflake/ml/experiment/callback.py

@@ -0,0 +1,121 @@
+import json
+from typing import TYPE_CHECKING, Any, Optional, Union
+from warnings import warn
+
+import lightgbm as lgb
+import xgboost as xgb
+
+from snowflake.ml.model.model_signature import ModelSignature
+
+if TYPE_CHECKING:
+    from snowflake.ml.experiment.experiment_tracking import ExperimentTracking
+
+
+class SnowflakeXgboostCallback(xgb.callback.TrainingCallback):
+    def __init__(
+        self,
+        experiment_tracking: "ExperimentTracking",
+        log_model: bool = True,
+        log_metrics: bool = True,
+        log_params: bool = True,
+        model_name: Optional[str] = None,
+        model_signature: Optional[ModelSignature] = None,
+    ) -> None:
+        self._experiment_tracking = experiment_tracking
+        self.log_model = log_model
+        self.log_metrics = log_metrics
+        self.log_params = log_params
+        self.model_name = model_name
+        self.model_signature = model_signature
+
+    def before_training(self, model: xgb.Booster) -> xgb.Booster:
+        def _flatten_nested_params(params: Union[list[Any], dict[str, Any]], prefix: str = "") -> dict[str, Any]:
+            flat_params = {}
+            items = params.items() if isinstance(params, dict) else enumerate(params)
+            for key, value in items:
+                new_prefix = f"{prefix}.{key}" if prefix else str(key)
+                if isinstance(value, (dict, list)):
+                    flat_params.update(_flatten_nested_params(value, new_prefix))
+                else:
+                    flat_params[new_prefix] = value
+            return flat_params
+
+        if self.log_params:
+            params = json.loads(model.save_config())
+            self._experiment_tracking.log_params(_flatten_nested_params(params))
+
+        return model
+
+    def after_iteration(self, model: Any, epoch: int, evals_log: dict[str, dict[str, Any]]) -> bool:
+        if self.log_metrics:
+            for dataset_name, metrics in evals_log.items():
+                for metric_name, log in metrics.items():
+                    metric_key = dataset_name + ":" + metric_name
+                    self._experiment_tracking.log_metric(key=metric_key, value=log[-1], step=epoch)
+
+        return False
+
+    def after_training(self, model: xgb.Booster) -> xgb.Booster:
+        if self.log_model:
+            if not self.model_signature:
+                warn(
+                    "Model will not be logged because model signature is missing. "
+                    "To autolog the model, please specify `model_signature` when constructing SnowflakeXgboostCallback."
+                )
+                return model
+
+            model_name = self.model_name or self._experiment_tracking._get_or_set_experiment().name + "_model"
+            self._experiment_tracking.log_model(  # type: ignore[call-arg]
+                model=model,
+                model_name=model_name,
+                signatures={"predict": self.model_signature},
+            )
+
+        return model
+
+
+class SnowflakeLightgbmCallback(lgb.callback._RecordEvaluationCallback):
+    def __init__(
+        self,
+        experiment_tracking: "ExperimentTracking",
+        log_model: bool = True,
+        log_metrics: bool = True,
+        log_params: bool = True,
+        model_name: Optional[str] = None,
+        model_signature: Optional[ModelSignature] = None,
+    ) -> None:
+        self._experiment_tracking = experiment_tracking
+        self.log_model = log_model
+        self.log_metrics = log_metrics
+        self.log_params = log_params
+        self.model_name = model_name
+        self.model_signature = model_signature
+
+        super().__init__(eval_result={})
+
+    def __call__(self, env: lgb.callback.CallbackEnv) -> None:
+        if self.log_params:
+            if env.iteration == env.begin_iteration:  # Log params only at the first iteration
+                self._experiment_tracking.log_params(env.params)
+
+        if self.log_metrics:
+            super().__call__(env)
+            for dataset_name, metrics in self.eval_result.items():
+                for metric_name, log in metrics.items():
+                    metric_key = dataset_name + ":" + metric_name
+                    self._experiment_tracking.log_metric(key=metric_key, value=log[-1], step=env.iteration)
+
+        if self.log_model:
+            if env.iteration == env.end_iteration - 1:  # Log model only at the last iteration
+                if self.model_signature:
+                    model_name = self.model_name or self._experiment_tracking._get_or_set_experiment().name + "_model"
+                    self._experiment_tracking.log_model(  # type: ignore[call-arg]
+                        model=env.model,
+                        model_name=model_name,
+                        signatures={"predict": self.model_signature},
+                    )
+                else:
+                    warn(
+                        "Model will not be logged because model signature is missing. To autolog the model, "
+                        "please specify `model_signature` when constructing SnowflakeLightgbmCallback."
+                    )

{snowflake_ml_python-1.9.1 → snowflake_ml_python-1.9.2}/snowflake/ml/jobs/_utils/constants.py

@@ -6,10 +6,23 @@ DEFAULT_CONTAINER_NAME = "main"
 PAYLOAD_DIR_ENV_VAR = "MLRS_PAYLOAD_DIR"
 RESULT_PATH_ENV_VAR = "MLRS_RESULT_PATH"
 MIN_INSTANCES_ENV_VAR = "MLRS_MIN_INSTANCES"
+TARGET_INSTANCES_ENV_VAR = "SNOWFLAKE_JOBS_COUNT"
 RUNTIME_IMAGE_TAG_ENV_VAR = "MLRS_CONTAINER_IMAGE_TAG"
 MEMORY_VOLUME_NAME = "dshm"
 STAGE_VOLUME_NAME = "stage-volume"
-STAGE_VOLUME_MOUNT_PATH = "/mnt/app"
+# Base mount path
+STAGE_VOLUME_MOUNT_PATH = "/mnt/job_stage"
+
+# Stage subdirectory paths
+APP_STAGE_SUBPATH = "app"
+SYSTEM_STAGE_SUBPATH = "system"
+OUTPUT_STAGE_SUBPATH = "output"
+
+# Complete mount paths (automatically generated from base + subpath)
+APP_MOUNT_PATH = f"{STAGE_VOLUME_MOUNT_PATH}/{APP_STAGE_SUBPATH}"
+SYSTEM_MOUNT_PATH = f"{STAGE_VOLUME_MOUNT_PATH}/{SYSTEM_STAGE_SUBPATH}"
+OUTPUT_MOUNT_PATH = f"{STAGE_VOLUME_MOUNT_PATH}/{OUTPUT_STAGE_SUBPATH}"
+
 
 # Default container image information
 DEFAULT_IMAGE_REPO = "/snowflake/images/snowflake_images"
@@ -46,9 +59,7 @@ ENABLE_HEALTH_CHECKS = "false"
 JOB_POLL_INITIAL_DELAY_SECONDS = 0.1
 JOB_POLL_MAX_DELAY_SECONDS = 30
 
-# Magic attributes
-IS_MLJOB_REMOTE_ATTR = "_is_mljob_remote_callable"
-RESULT_PATH_DEFAULT_VALUE = "mljob_result.pkl"
+RESULT_PATH_DEFAULT_VALUE = f"{OUTPUT_MOUNT_PATH}/mljob_result.pkl"
 
 # Log start and end messages
 LOG_START_MSG = "--------------------------------\nML job started\n--------------------------------"