Project import generated by Copybara. (#178)

GitOrigin-RevId: f10ddb99e98b5c15b4b34ebcd31a289218c43a3c

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

Author: sfc-gh-anavalos

.bazelrc

@@ -31,13 +31,13 @@ build:py3.12 --repo_env=BAZEL_CONDA_PYTHON_VERSION=3.12
 build:build --config=_build
 
 # Config to sync files
-run:pre_build --config=_build --config=py3.9
+run:pre_build --config=_build --config=py3.10
 
 # Config to run type check
-build:typecheck --aspects @rules_mypy//:mypy.bzl%mypy_aspect --output_groups=mypy --config=_all --config=py3.9
+build:typecheck --aspects @rules_mypy//:mypy.bzl%mypy_aspect --output_groups=mypy --config=_all --config=py3.10
 
 # Config to build the doc
-build:docs --config=_all --config=py3.9
+build:docs --config=_all --config=py3.10
 
 # Public the extended setting
 

CHANGELOG.md

@@ -1,5 +1,15 @@
 # Release History
 
+## 1.14.0
+
+### Bug Fixes
+
+### Behavior Changes
+
+### New Features
+
+* ML Job: The `additional_payloads` argument is now **deprecated** in favor of `imports`.
+
 ## 1.13.0
 
 ### Bug Fixes

bazel/environments/fetch_conda_env_config.bzl

@@ -3,7 +3,7 @@ load("//bazel/platforms:optional_dependency_groups.bzl", "OPTIONAL_DEPENDENCY_GR
 def _fetch_conda_env_config_impl(rctx):
     # read the particular environment variable we are interested in
     env_name = rctx.os.environ.get("BAZEL_CONDA_ENV_NAME", "core").lower()
-    python_ver = rctx.os.environ.get("BAZEL_CONDA_PYTHON_VERSION", "3.9").lower()
+    python_ver = rctx.os.environ.get("BAZEL_CONDA_PYTHON_VERSION", "3.10").lower()
 
     # necessary to create empty BUILD file for this rule
     # which will be located somewhere in the Bazel build files

bazel/requirements/templates/bazelrc.tpl

@@ -28,13 +28,13 @@ build:py3.12 --repo_env=BAZEL_CONDA_PYTHON_VERSION=3.12
 build:build --config=_build
 
 # Config to sync files
-run:pre_build --config=_build --config=py3.9
+run:pre_build --config=_build --config=py3.10
 
 # Config to run type check
-build:typecheck --aspects @rules_mypy//:mypy.bzl%mypy_aspect --output_groups=mypy --config=_all --config=py3.9
+build:typecheck --aspects @rules_mypy//:mypy.bzl%mypy_aspect --output_groups=mypy --config=_all --config=py3.10
 
 # Config to build the doc
-build:docs --config=_all --config=py3.9
+build:docs --config=_all --config=py3.10
 
 # Public the extended setting
 

ci/build_and_run_tests.sh

@@ -42,7 +42,7 @@ WITH_SNOWPARK=false
 WITH_SPCS_IMAGE=false
 RUN_GRYPE=false
 MODE="continuous_run"
-PYTHON_VERSION=3.9
+PYTHON_VERSION=3.10
 PYTHON_ENABLE_SCRIPT="bin/activate"
 SNOWML_DIR="snowml"
 SNOWPARK_DIR="snowpark-python"

ci/conda_recipe/meta.yaml

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

snowflake/ml/jobs/_utils/constants.py

@@ -25,7 +25,7 @@
 DEFAULT_IMAGE_REPO = "/snowflake/images/snowflake_images"
 DEFAULT_IMAGE_CPU = "st_plat/runtime/x86/runtime_image/snowbooks"
 DEFAULT_IMAGE_GPU = "st_plat/runtime/x86/generic_gpu/runtime_image/snowbooks"
-DEFAULT_IMAGE_TAG = "1.6.2"
+DEFAULT_IMAGE_TAG = "1.8.0"
 DEFAULT_ENTRYPOINT_PATH = "func.py"
 
 # Percent of container memory to allocate for /dev/shm volume

snowflake/ml/jobs/_utils/scripts/mljob_launcher.py

@@ -234,12 +234,6 @@ def run_script(script_path: str, *script_args: Any, main_func: Optional[str] = N
     if payload_dir and payload_dir not in sys.path:
         sys.path.insert(0, payload_dir)
 
-    # Create a Snowpark session before running the script
-    # Session can be retrieved from using snowflake.snowpark.context.get_active_session()
-    config = SnowflakeLoginOptions()
-    config["client_session_keep_alive"] = "True"
-    session = Session.builder.configs(config).create()  # noqa: F841
-
     try:
 
         if main_func:
@@ -266,7 +260,6 @@ def run_script(script_path: str, *script_args: Any, main_func: Optional[str] = N
     finally:
         # Restore original sys.argv
         sys.argv = original_argv
-        session.close()
 
 
 def main(script_path: str, *script_args: Any, script_main_func: Optional[str] = None) -> ExecutionResult:
@@ -297,6 +290,12 @@ def main(script_path: str, *script_args: Any, script_main_func: Optional[str] =
     except ModuleNotFoundError:
         warnings.warn("Ray is not installed, skipping Ray initialization", ImportWarning, stacklevel=1)
 
+    # Create a Snowpark session before starting
+    # Session can be retrieved from using snowflake.snowpark.context.get_active_session()
+    config = SnowflakeLoginOptions()
+    config["client_session_keep_alive"] = "True"
+    session = Session.builder.configs(config).create()  # noqa: F841
+
     try:
         # Wait for minimum required instances if specified
         min_instances_str = os.environ.get(MIN_INSTANCES_ENV_VAR) or "1"
@@ -352,6 +351,9 @@ def main(script_path: str, *script_args: Any, script_main_func: Optional[str] =
                 f"Failed to serialize JSON result to {result_json_path}: {json_exc}", RuntimeWarning, stacklevel=1
             )
 
+        # Close the session after serializing the result
+        session.close()
+
 
 if __name__ == "__main__":
     # Parse command line arguments

snowflake/ml/jobs/job.py

@@ -83,6 +83,8 @@ def _service_spec(self) -> dict[str, Any]:
     def _container_spec(self) -> dict[str, Any]:
         """Get the job's main container spec."""
         containers = self._service_spec["spec"]["containers"]
+        if len(containers) == 1:
+            return cast(dict[str, Any], containers[0])
         try:
             container_spec = next(c for c in containers if c["name"] == constants.DEFAULT_CONTAINER_NAME)
         except StopIteration:
@@ -163,7 +165,7 @@ def get_logs(
         Returns:
             The job's execution logs.
         """
-        logs = _get_logs(self._session, self.id, limit, instance_id, verbose)
+        logs = _get_logs(self._session, self.id, limit, instance_id, self._container_spec["name"], verbose)
         assert isinstance(logs, str)  # mypy
         if as_list:
             return logs.splitlines()
@@ -281,7 +283,12 @@ def _get_service_spec(session: snowpark.Session, job_id: str) -> dict[str, Any]:
 
 @telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["job_id", "limit", "instance_id"])
 def _get_logs(
-    session: snowpark.Session, job_id: str, limit: int = -1, instance_id: Optional[int] = None, verbose: bool = True
+    session: snowpark.Session,
+    job_id: str,
+    limit: int = -1,
+    instance_id: Optional[int] = None,
+    container_name: str = constants.DEFAULT_CONTAINER_NAME,
+    verbose: bool = True,
 ) -> str:
     """
     Retrieve the job's execution logs.
@@ -291,6 +298,7 @@ def _get_logs(
         limit: The maximum number of lines to return. Negative values are treated as no limit.
         session: The Snowpark session to use. If none specified, uses active session.
         instance_id: Optional instance ID to get logs from a specific instance.
+        container_name: The container name to get logs from a specific container.
         verbose: Whether to return the full log or just the portion between START and END messages.
 
     Returns:
@@ -311,7 +319,7 @@ def _get_logs(
     params: list[Any] = [
         job_id,
         0 if instance_id is None else instance_id,
-        constants.DEFAULT_CONTAINER_NAME,
+        container_name,
     ]
     if limit > 0:
         params.append(limit)
@@ -337,7 +345,7 @@ def _get_logs(
                     job_id,
                     limit=limit,
                     instance_id=instance_id if instance_id else 0,
-                    container_name=constants.DEFAULT_CONTAINER_NAME,
+                    container_name=container_name,
                 )
                 full_log = os.linesep.join(row[0] for row in logs)
 

snowflake/ml/jobs/jobs_test.py

@@ -10,6 +10,13 @@
 from snowflake.snowpark import exceptions as sp_exceptions
 from snowflake.snowpark.row import Row
 
+SERVICE_SPEC = """
+spec:
+  containers:
+    - name: main
+      image: test-image
+"""
+
 
 class JobTest(parameterized.TestCase):
     @parameterized.named_parameters(  # type: ignore[misc]
@@ -83,7 +90,7 @@ def test_get_logs_negative(self) -> None:
 
         def sql_side_effect(session: snowpark.Session, query_str: str, *args: Any, **kwargs: Any) -> Any:
             if query_str.startswith("DESCRIBE SERVICE IDENTIFIER"):
-                return [Row(target_instances=2)]
+                return [Row(target_instances=2, spec=SERVICE_SPEC)]
             else:
                 raise sp_exceptions.SnowparkSQLException("Waiting to start, Container Status: PENDING")
 
@@ -97,7 +104,7 @@ def test_get_logs_from_event_table(self) -> None:
         def sql_side_effect(session: snowpark.Session, query_str: str, *args: Any, **kwargs: Any) -> Any:
             if query_str.startswith("DESCRIBE SERVICE IDENTIFIER"):
                 return [
-                    Row(target_instances=2),
+                    Row(target_instances=2, spec=SERVICE_SPEC),
                 ]
             elif query_str.startswith("SELECT VALUE FROM "):
                 return [