Project import generated by Copybara. (#170)

GitOrigin-RevId: 964e61b5d23d49642fd754c11e3f0b2707e79ccd

Co-authored-by: Snowflake Authors <[email protected]>

Author: sfc-gh-anavalos

.pre-commit-config.yaml

@@ -6,7 +6,7 @@ repos:
       - id: check-py-test-feature-tags
         name: Check py_test feature area tags
         description: Ensure all py_test targets have feature area tags
-        entry: python bazel/check_feature_tags.py --precommit
+        entry: python3 bazel/check_feature_tags.py --precommit
         language: system
         files: BUILD\.bazel$
         stages:

CHANGELOG.md

@@ -1,7 +1,7 @@
 #!/bin/bash
 # DESCRIPTION: Utility Shell script to run bazel action for snowml repository
 #
-# RunBazelAction.sh <test|coverage> [-b <bazel_path>] [-m merge_gate|continuous_run|quarantined|local_unittest|local_all] [-t <target>] [-c <path_to_coverage_report>]
+# RunBazelAction.sh <test|coverage> [-b <bazel_path>] [-m merge_gate|continuous_run|quarantined|local_unittest|local_all] [-t <target>] [-c <path_to_coverage_report>] [--tags <tags>]
 #
 # Args:
 #   action: bazel action, choose from test and coverage
@@ -17,6 +17,7 @@
 #   -t: specify the target for local_unit and local_all mode
 #   -c: specify the path to the coverage report dat file.
 #   -e: specify the environment, used to determine.
+#   --tags: specify bazel test tag filters (e.g., "feature:jobs,feature:data")
 #
 
 set -o pipefail
@@ -28,13 +29,21 @@ mode="continuous_run"
 target=""
 SF_ENV="prod3"
 WITH_SPCS_IMAGE=false
+TAG_FILTERS=""
 PROG=$0
 
 action=$1 && shift
 
 help() {
     local exit_code=$1
-    echo "Usage: ${PROG} <test|coverage> [-b <bazel_path>] [-m merge_gate|continuous_run|quarantined|local_unittest|local_all|perf] [-e <snowflake_env>] [--with-spcs-image]"
+    echo "Usage: ${PROG} <test|coverage> [-b <bazel_path>] [-m merge_gate|continuous_run|quarantined|local_unittest|local_all|perf] [-e <snowflake_env>] [--tags <tags>] [--with-spcs-image]"
+    echo ""
+    echo "Options:"
+    echo "  --tags <tags>       Specify bazel tag filters (comma-separated)"
+    echo ""
+    echo "Examples:"
+    echo "  ${PROG} test --tags 'feature:jobs'"
+    echo "  ${PROG} test --tags 'feature:jobs,feature:data'"
     exit "${exit_code}"
 }
 
@@ -72,6 +81,10 @@ while (($#)); do
         shift
         SF_ENV=$1
         ;;
+    --tags)
+        shift
+        TAG_FILTERS="$1"
+        ;;
     --with-spcs-image)
         WITH_SPCS_IMAGE=true
         ;;
@@ -98,13 +111,20 @@ action_env=()
 if [[ "${WITH_SPCS_IMAGE}" = true ]]; then
     export SKIP_GRYPE=true
     source model_container_services_deployment/ci/build_and_push_images.sh
-    action_env=("--action_env=BUILDER_IMAGE_PATH=${BUILDER_IMAGE_PATH}" "--action_env=BASE_CPU_IMAGE_PATH=${BASE_CPU_IMAGE_PATH}" "--action_env=BASE_GPU_IMAGE_PATH=${BASE_GPU_IMAGE_PATH}" "--action_env=IMAGE_BUILD_SIDECAR_CPU_PATH=${IMAGE_BUILD_SIDECAR_CPU_PATH}" "--action_env=IMAGE_BUILD_SIDECAR_GPU_PATH=${IMAGE_BUILD_SIDECAR_GPU_PATH}" "--action_env=PROXY_IMAGE_PATH=${PROXY_IMAGE_PATH}")
+    action_env=("--action_env=BUILDER_IMAGE_PATH=${BUILDER_IMAGE_PATH}" "--action_env=BASE_CPU_IMAGE_PATH=${BASE_CPU_IMAGE_PATH}" "--action_env=BASE_GPU_IMAGE_PATH=${BASE_GPU_IMAGE_PATH}" "--action_env=IMAGE_BUILD_SIDECAR_CPU_PATH=${IMAGE_BUILD_SIDECAR_CPU_PATH}" "--action_env=IMAGE_BUILD_SIDECAR_GPU_PATH=${IMAGE_BUILD_SIDECAR_GPU_PATH}" "--action_env=PROXY_IMAGE_PATH=${PROXY_IMAGE_PATH}" "--action_env=VLLM_IMAGE_PATH=${VLLM_IMAGE_PATH}")
 fi
 
 working_dir=$(mktemp -d "/tmp/tmp_XXXXX")
 trap 'rm -rf "${working_dir}"' EXIT
 
-tag_filter="--test_tag_filters="
+# Set up tag filtering
+if [[ -n "${TAG_FILTERS:-}" ]]; then
+    tag_filter="--test_tag_filters=${TAG_FILTERS}"
+    echo "Running with tag filters: ${TAG_FILTERS}"
+else
+    tag_filter="--test_tag_filters="
+fi
+
 cache_test_results="--cache_test_results=no"
 
 case "${mode}" in

ci/RunBazelAction.sh

@@ -17,7 +17,7 @@ build:
   noarch: python
 package:
   name: snowflake-ml-python
-  version: 1.9.2
+  version: 1.10.0
 requirements:
   build:
     - python

ci/conda_recipe/meta.yaml

@@ -9,6 +9,15 @@
 
 import platformdirs
 
+# Module-level logger for operational messages that should appear on console
+stdout_handler = logging.StreamHandler(sys.stdout)
+stdout_handler.setFormatter(logging.Formatter("%(message)s"))
+
+console_logger = logging.getLogger(__name__)
+console_logger.addHandler(stdout_handler)
+console_logger.setLevel(logging.INFO)
+console_logger.propagate = False
+
 
 class LogColor(enum.Enum):
     GREY = "\x1b[38;20m"
@@ -109,49 +118,54 @@ 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)
+    parent_logger.setLevel(logging.DEBUG)
+    parent_logger.propagate = False
 
-    # 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}")
+                console_logger.info(f"create_service 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.")
+                console_logger.warning(f"Could not create log file at {log_file_path}: {e}.")
         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.")
+            console_logger.warning("No writable location found for create_service log file.")
+
+        if logging.getLogger().level > logging.INFO:
+            console_logger.info(
+                "To see logs in console, set log level to INFO: logging.getLogger().setLevel(logging.INFO)"
+            )
 
     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)
+    root_logger = logging.getLogger()
 
     # 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
 
+        if root_logger.level <= logging.INFO:
+            handler = logging.StreamHandler(sys.stdout)
+            handler.setFormatter(CustomFormatter(info_color))
+            logger.addHandler(handler)
+    else:
+        # No operation_id - add console handler only if user wants verbose logging
+        if root_logger.level <= logging.INFO and not logger.handlers:
+            handler = logging.StreamHandler(sys.stdout)
+            handler.setFormatter(CustomFormatter(info_color))
+            logger.addHandler(handler)
+
     return logger
 
 

snowflake/ml/_internal/utils/service_logger.py

@@ -17,11 +17,17 @@ def setUp(self) -> None:
         self.test_logger_name = "test_logger"
         self.test_color = service_logger.LogColor.BLUE
 
+        # Store original root logger level
+        self._original_root_level = logging.getLogger().level
+
         # Clear any existing loggers to avoid test interference
         self._clear_all_loggers()
 
     def tearDown(self) -> None:
         """Clean up test fixtures."""
+        # Restore original root logger level
+        logging.getLogger().setLevel(self._original_root_level)
+
         # Clear loggers created during tests
         self._clear_all_loggers()
 
@@ -45,9 +51,12 @@ def _clear_all_loggers(self) -> None:
 
     def test_get_logger_without_operation_id(self) -> None:
         """Test that get_logger works without operation_id (backward compatibility)."""
+        # Set root logger to INFO level to ensure handler is added
+        logging.getLogger().setLevel(logging.INFO)
+
         logger = service_logger.get_logger(self.test_logger_name, self.test_color)
 
-        # Should have one handler (StreamHandler)
+        # Should have one handler (StreamHandler) when in verbose mode
         self.assertEqual(len(logger.handlers), 1)
         self.assertIsInstance(logger.handlers[0], logging.StreamHandler)
 
@@ -58,11 +67,24 @@ def test_get_logger_without_operation_id(self) -> None:
         # Should use custom formatter
         self.assertIsInstance(logger.handlers[0].formatter, service_logger.CustomFormatter)
 
+    def test_get_logger_without_operation_id_quiet_mode(self) -> None:
+        """Test that get_logger without operation_id doesn't add handler in quiet mode."""
+        # Set root logger to WARNING level (quiet mode)
+        logging.getLogger().setLevel(logging.WARNING)
+
+        logger = service_logger.get_logger(self.test_logger_name, self.test_color)
+
+        # Should have no handlers in quiet mode
+        self.assertEqual(len(logger.handlers), 0)
+
     def test_get_logger_with_operation_id(self) -> None:
-        """Test that get_logger creates parent logger with operation_id."""
+        """Test that get_logger creates parent logger with operation_id in verbose mode."""
+        # Set root logger to INFO level to ensure handler is added
+        logging.getLogger().setLevel(logging.INFO)
+
         logger = service_logger.get_logger(self.test_logger_name, self.test_color, self.test_operation_id)
 
-        # Should have one handler (StreamHandler)
+        # Should have one handler (StreamHandler) in verbose mode
         self.assertEqual(len(logger.handlers), 1)
         self.assertIsInstance(logger.handlers[0], logging.StreamHandler)
 
@@ -74,6 +96,24 @@ def test_get_logger_with_operation_id(self) -> None:
         # Should propagate to parent
         self.assertTrue(logger.propagate)
 
+    def test_get_logger_with_operation_id_quiet_mode(self) -> None:
+        """Test that get_logger with operation_id doesn't add console handler in quiet mode."""
+        # Set root logger to WARNING level (quiet mode)
+        logging.getLogger().setLevel(logging.WARNING)
+
+        logger = service_logger.get_logger(self.test_logger_name, self.test_color, self.test_operation_id)
+
+        # Should have no console handlers in quiet mode (only file via propagation)
+        self.assertEqual(len(logger.handlers), 0)
+
+        # Should still have parent logger for file output
+        self.assertIsNotNone(logger.parent)
+        if logger.parent:
+            self.assertEqual(logger.parent.name, f"snowflake_ml_operation_{self.test_operation_id}")
+
+        # Should still propagate to parent for file logging
+        self.assertTrue(logger.propagate)
+
     def test_parent_logger_creation_with_file(self) -> None:
         """Test that parent logger is created with FileHandler when filesystem is writable."""
         with mock.patch("snowflake.ml._internal.utils.service_logger._get_log_file_path") as mock_path:

snowflake/ml/_internal/utils/service_logger_test.py

@@ -2,29 +2,6 @@ load("//bazel:py_rules.bzl", "py_library", "py_package", "py_test")
 
 package(default_visibility = ["//visibility:public"])
 
-py_library(
-    name = "callback",
-    srcs = ["callback.py"],
-    deps = [
-        "//snowflake/ml/experiment:experiment_tracking",
-        "//snowflake/ml/model:model_signature",
-    ],
-)
-
-py_test(
-    name = "callback_test",
-    srcs = ["callback_test.py"],
-    optional_dependencies = [
-        "lightgbm",
-    ],
-    tags = ["feature:observability"],
-    deps = [
-        ":callback",
-        "//snowflake/ml/experiment:experiment_tracking",
-        "//snowflake/ml/model:model_signature",
-    ],
-)
-
 py_library(
     name = "experiment_info",
     srcs = ["_experiment_info.py"],
@@ -38,7 +15,10 @@ py_test(
     srcs = ["_experiment_info_test.py"],
     main = "_experiment_info_test.py",
     tags = ["feature:observability"],
-    deps = [":experiment_info"],
+    deps = [
+        ":experiment_info",
+        "//snowflake/ml/test_utils:mock_progress",
+    ],
 )
 
 py_library(
@@ -62,12 +42,24 @@ py_test(
     ],
 )
 
+py_library(
+    name = "utils",
+    srcs = ["utils.py"],
+)
+
+py_test(
+    name = "utils_test",
+    srcs = ["utils_test.py"],
+    tags = ["feature:observability"],
+    deps = [":utils"],
+)
+
 py_library(
     name = "experiment",
     srcs = ["__init__.py"],
     deps = [
-        ":callback",
         ":experiment_tracking",
+        "//snowflake/ml/experiment/callback",
     ],
 )
 

snowflake/ml/experiment/BUILD.bazel

@@ -4,6 +4,7 @@
 
 from snowflake.ml.experiment import _experiment_info as experiment_info
 from snowflake.ml.registry._manager.model_manager import ModelManager
+from snowflake.ml.test_utils.mock_progress import create_mock_progress_status
 
 
 class ExperimentInfoPatcherTest(absltest.TestCase):
@@ -89,6 +90,7 @@ def test_preserves_original_method_signature_and_return_value(self) -> None:
         experiment_info.ExperimentInfoPatcher._original_log_model = mock_original
 
         mock_manager = mock.MagicMock(spec=ModelManager)
+        mock_progress_status = create_mock_progress_status()
 
         with experiment_info.ExperimentInfoPatcher(self.experiment_info):
             # Call with various arguments
@@ -99,6 +101,7 @@ def test_preserves_original_method_signature_and_return_value(self) -> None:
                 version_name="v1.0",
                 comment="test comment",
                 metrics={"accuracy": 0.95},
+                progress_status=mock_progress_status,
             )
 
             # Verify all arguments are passed through with experiment_info added
@@ -109,6 +112,7 @@ def test_preserves_original_method_signature_and_return_value(self) -> None:
                 version_name="v1.0",
                 comment="test comment",
                 metrics={"accuracy": 0.95},
+                progress_status=mock_progress_status,
                 experiment_info=self.experiment_info,
             )