snowflake-ml-python 1.13.0__py3-none-any.whl → 1.15.0__py3-none-any.whl

snowflake/ml/_internal/platform_capabilities.py

@@ -1,8 +1,8 @@
 import json
+import logging
 from contextlib import contextmanager
 from typing import Any, Optional
 
-from absl import logging
 from packaging import version
 
 from snowflake.ml import version as snowml_version
@@ -13,6 +13,8 @@ from snowflake.snowpark import (
     session as snowpark_session,
 )
 
+logger = logging.getLogger(__name__)
+
 LIVE_COMMIT_PARAMETER = "ENABLE_LIVE_VERSION_IN_SDK"
 INLINE_DEPLOYMENT_SPEC_PARAMETER = "ENABLE_INLINE_DEPLOYMENT_SPEC_FROM_CLIENT_VERSION"
 
@@ -60,12 +62,12 @@ class PlatformCapabilities:
     @classmethod  # type: ignore[arg-type]
     @contextmanager
     def mock_features(cls, features: dict[str, Any] = _dummy_features) -> None:  # type: ignore[misc]
-        logging.debug(f"Setting mock features: {features}")
+        logger.debug(f"Setting mock features: {features}")
         cls.set_mock_features(features)
         try:
             yield
         finally:
-            logging.debug(f"Clearing mock features: {features}")
+            logger.debug(f"Clearing mock features: {features}")
             cls.clear_mock_features()
 
     def is_inlined_deployment_spec_enabled(self) -> bool:
@@ -98,7 +100,7 @@ class PlatformCapabilities:
                         error_code=error_codes.INTERNAL_SNOWML_ERROR, original_exception=RuntimeError(message)
                     )
         except snowpark_exceptions.SnowparkSQLException as e:
-            logging.debug(f"Failed to retrieve platform capabilities: {e}")
+            logger.debug(f"Failed to retrieve platform capabilities: {e}")
             # This can happen is server side is older than 9.2. That is fine.
         return {}
 
@@ -144,7 +146,7 @@ class PlatformCapabilities:
 
         value = self.features.get(feature_name)
         if value is None:
-            logging.debug(f"Feature {feature_name} not found, returning large version number")
+            logger.debug(f"Feature {feature_name} not found, returning large version number")
             return large_version
 
         try:
@@ -152,7 +154,7 @@ class PlatformCapabilities:
             version_str = str(value)
             return version.Version(version_str)
         except (version.InvalidVersion, ValueError, TypeError) as e:
-            logging.debug(
+            logger.debug(
                 f"Failed to parse version from feature {feature_name} with value '{value}': {e}. "
                 f"Returning large version number"
             )
@@ -171,7 +173,7 @@ class PlatformCapabilities:
         feature_version = self._get_version_feature(feature_name)
 
         result = current_version >= feature_version
-        logging.debug(
+        logger.debug(
             f"Version comparison for feature {feature_name}: "
             f"current={current_version}, feature={feature_version}, enabled={result}"
         )

snowflake/ml/_internal/utils/connection_params.py

@@ -1,11 +1,13 @@
 import configparser
+import logging
 import os
 from typing import Optional, Union
 
-from absl import logging
 from cryptography.hazmat import backends
 from cryptography.hazmat.primitives import serialization
 
+logger = logging.getLogger(__name__)
+
 _DEFAULT_CONNECTION_FILE = "~/.snowsql/config"
 
 
@@ -106,7 +108,7 @@ def _load_from_snowsql_config_file(connection_name: str, login_file: str = "") -
     """Loads the dictionary from snowsql config file."""
     snowsql_config_file = login_file if login_file else os.path.expanduser(_DEFAULT_CONNECTION_FILE)
     if not os.path.exists(snowsql_config_file):
-        logging.error(f"Connection name given but snowsql config file is not found at: {snowsql_config_file}")
+        logger.error(f"Connection name given but snowsql config file is not found at: {snowsql_config_file}")
         raise Exception("Snowflake SnowSQL config not found.")
 
     config = configparser.ConfigParser(inline_comment_prefixes="#")
@@ -122,7 +124,7 @@ def _load_from_snowsql_config_file(connection_name: str, login_file: str = "") -
         # See https://docs.snowflake.com/en/user-guide/snowsql-start.html#configuring-default-connection-settings
         connection_name = "connections"
 
-    logging.info(f"Reading {snowsql_config_file} for connection parameters defined as {connection_name}")
+    logger.info(f"Reading {snowsql_config_file} for connection parameters defined as {connection_name}")
     config.read(snowsql_config_file)
     conn_params = dict(config[connection_name])
     # Remap names to appropriate args in Python Connector API

snowflake/ml/_internal/utils/jwt_generator.py

@@ -110,15 +110,16 @@ class JWTGenerator:
         }
 
         # Regenerate the actual token
-        token = jwt.encode(payload, key=self.private_key, algorithm=JWTGenerator.ALGORITHM)
+        token = jwt.encode(payload, key=self.private_key, algorithm=JWTGenerator.ALGORITHM)  # type: ignore[arg-type]
         # If you are using a version of PyJWT prior to 2.0, jwt.encode returns a byte string instead of a string.
         # If the token is a byte string, convert it to a string.
         if isinstance(token, bytes):
             token = token.decode("utf-8")
         self.token = token
+        public_key = self.private_key.public_key()
         logger.info(
             "Generated a JWT with the following payload: %s",
-            jwt.decode(self.token, key=self.private_key.public_key(), algorithms=[JWTGenerator.ALGORITHM]),
+            jwt.decode(self.token, key=public_key, algorithms=[JWTGenerator.ALGORITHM]),  # type: ignore[arg-type]
         )
 
         return token

snowflake/ml/_internal/utils/temp_file_utils.py

@@ -1,10 +1,9 @@
+import logging
 import os
 import shutil
 import tempfile
 from typing import Iterable, Union
 
-from absl.logging import logging
-
 logger = logging.getLogger(__name__)
 
 

snowflake/ml/experiment/_client/experiment_tracking_sql_client.py

@@ -76,17 +76,30 @@ class ExperimentTrackingSQLClient(_base._BaseSQLClient):
             self._session, f"ALTER EXPERIMENT {experiment_fqn} DROP RUN {run_name}"
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
-    def modify_run(
+    def modify_run_add_metrics(
         self,
         *,
         experiment_name: sql_identifier.SqlIdentifier,
         run_name: sql_identifier.SqlIdentifier,
-        run_metadata: str,
+        metrics: str,
     ) -> None:
         experiment_fqn = self.fully_qualified_object_name(self._database_name, self._schema_name, experiment_name)
         query_result_checker.SqlResultValidator(
             self._session,
-            f"ALTER EXPERIMENT {experiment_fqn} MODIFY RUN {run_name} SET METADATA=$${run_metadata}$$",
+            f"ALTER EXPERIMENT {experiment_fqn} MODIFY RUN {run_name} ADD METRICS=$${metrics}$$",
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def modify_run_add_params(
+        self,
+        *,
+        experiment_name: sql_identifier.SqlIdentifier,
+        run_name: sql_identifier.SqlIdentifier,
+        params: str,
+    ) -> None:
+        experiment_fqn = self.fully_qualified_object_name(self._database_name, self._schema_name, experiment_name)
+        query_result_checker.SqlResultValidator(
+            self._session,
+            f"ALTER EXPERIMENT {experiment_fqn} MODIFY RUN {run_name} ADD PARAMETERS=$${params}$$",
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
     def put_artifact(

snowflake/ml/experiment/_entities/__init__.py

@@ -1,4 +1,5 @@
 from snowflake.ml.experiment._entities.experiment import Experiment
 from snowflake.ml.experiment._entities.run import Run
+from snowflake.ml.experiment._entities.run_metadata import Metric, Param
 
-__all__ = ["Experiment", "Run"]
+__all__ = ["Experiment", "Run", "Metric", "Param"]

snowflake/ml/experiment/_entities/run.py

@@ -1,11 +1,8 @@
-import json
 import types
 from typing import TYPE_CHECKING, Optional
 
 from snowflake.ml._internal.utils import sql_identifier
 from snowflake.ml.experiment import _experiment_info as experiment_info
-from snowflake.ml.experiment._client import experiment_tracking_sql_client
-from snowflake.ml.experiment._entities import run_metadata
 
 if TYPE_CHECKING:
     from snowflake.ml.experiment import experiment_tracking
@@ -41,18 +38,6 @@ class Run:
         if self._experiment_tracking._run is self:
             self._experiment_tracking.end_run()
 
-    def _get_metadata(
-        self,
-    ) -> run_metadata.RunMetadata:
-        runs = self._experiment_tracking._sql_client.show_runs_in_experiment(
-            experiment_name=self.experiment_name, like=str(self.name)
-        )
-        if not runs:
-            raise RuntimeError(f"Run {self.name} not found in experiment {self.experiment_name}.")
-        return run_metadata.RunMetadata.from_dict(
-            json.loads(runs[0][experiment_tracking_sql_client.ExperimentTrackingSQLClient.RUN_METADATA_COL_NAME])
-        )
-
     def _get_experiment_info(self) -> experiment_info.ExperimentInfo:
         return experiment_info.ExperimentInfo(
             fully_qualified_name=self._experiment_tracking._sql_client.fully_qualified_object_name(

snowflake/ml/experiment/_entities/run_metadata.py

@@ -1,12 +1,4 @@
 import dataclasses
-import enum
-import typing
-
-
-class RunStatus(str, enum.Enum):
-    UNKNOWN = "UNKNOWN"
-    RUNNING = "RUNNING"
-    FINISHED = "FINISHED"
 
 
 @dataclasses.dataclass
@@ -15,54 +7,14 @@ class Metric:
     value: float
     step: int
 
+    def to_dict(self) -> dict:  # type: ignore[type-arg]
+        return dataclasses.asdict(self)
+
 
 @dataclasses.dataclass
 class Param:
     name: str
     value: str
 
-
-@dataclasses.dataclass
-class RunMetadata:
-    status: RunStatus
-    metrics: list[Metric]
-    parameters: list[Param]
-
-    @classmethod
-    def from_dict(
-        cls,
-        metadata: dict,  # type: ignore[type-arg]
-    ) -> "RunMetadata":
-        return RunMetadata(
-            status=RunStatus(metadata.get("status", RunStatus.UNKNOWN.value)),
-            metrics=[Metric(**m) for m in metadata.get("metrics", [])],
-            parameters=[Param(**p) for p in metadata.get("parameters", [])],
-        )
-
     def to_dict(self) -> dict:  # type: ignore[type-arg]
         return dataclasses.asdict(self)
-
-    def set_metric(
-        self,
-        key: str,
-        value: float,
-        step: int,
-    ) -> None:
-        for metric in self.metrics:
-            if metric.name == key and metric.step == step:
-                metric.value = value
-                break
-        else:
-            self.metrics.append(Metric(name=key, value=value, step=step))
-
-    def set_param(
-        self,
-        key: str,
-        value: typing.Any,
-    ) -> None:
-        for parameter in self.parameters:
-            if parameter.name == key:
-                parameter.value = str(value)
-                break
-        else:
-            self.parameters.append(Param(name=key, value=str(value)))

snowflake/ml/experiment/experiment_tracking.py

@@ -261,13 +261,13 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
             step: The step of the metrics. Defaults to 0.
         """
         run = self._get_or_start_run()
-        metadata = run._get_metadata()
+        metrics_list = []
         for key, value in metrics.items():
-            metadata.set_metric(key, value, step)
-        self._sql_client.modify_run(
+            metrics_list.append(entities.Metric(key, value, step))
+        self._sql_client.modify_run_add_metrics(
             experiment_name=run.experiment_name,
             run_name=run.name,
-            run_metadata=json.dumps(metadata.to_dict()),
+            metrics=json.dumps([metric.to_dict() for metric in metrics_list]),
         )
 
     def log_param(
@@ -296,13 +296,13 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
                 to string.
         """
         run = self._get_or_start_run()
-        metadata = run._get_metadata()
+        params_list = []
         for key, value in params.items():
-            metadata.set_param(key, value)
-        self._sql_client.modify_run(
+            params_list.append(entities.Param(key, str(value)))
+        self._sql_client.modify_run_add_params(
             experiment_name=run.experiment_name,
             run_name=run.name,
-            run_metadata=json.dumps(metadata.to_dict()),
+            params=json.dumps([param.to_dict() for param in params_list]),
         )
 
     def log_artifact(

snowflake/ml/jobs/_utils/constants.py

@@ -25,7 +25,7 @@ RESULT_PATH_DEFAULT_VALUE = f"{OUTPUT_STAGE_SUBPATH}/mljob_result.pkl"
 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 @@ class MLJob(Generic[T], SerializableSessionMixin):
     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 @@ class MLJob(Generic[T], SerializableSessionMixin):
         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/manager.py

@@ -232,6 +232,7 @@ def submit_file(
             enable_metrics (bool): Whether to enable metrics publishing for the job.
             query_warehouse (str): The query warehouse to use. Defaults to session warehouse.
             spec_overrides (dict): A dictionary of overrides for the service spec.
+            imports (list[Union[tuple[str, str], tuple[str]]]): A list of additional payloads used in the job.
 
     Returns:
         An object representing the submitted job.
@@ -286,6 +287,7 @@ def submit_directory(
             enable_metrics (bool): Whether to enable metrics publishing for the job.
             query_warehouse (str): The query warehouse to use. Defaults to session warehouse.
             spec_overrides (dict): A dictionary of overrides for the service spec.
+            imports (list[Union[tuple[str, str], tuple[str]]]): A list of additional payloads used in the job.
 
     Returns:
         An object representing the submitted job.
@@ -341,6 +343,7 @@ def submit_from_stage(
             enable_metrics (bool): Whether to enable metrics publishing for the job.
             query_warehouse (str): The query warehouse to use. Defaults to session warehouse.
             spec_overrides (dict): A dictionary of overrides for the service spec.
+            imports (list[Union[tuple[str, str], tuple[str]]]): A list of additional payloads used in the job.
 
     Returns:
         An object representing the submitted job.
@@ -404,6 +407,8 @@ def _submit_job(
         "num_instances",  # deprecated
         "target_instances",
         "min_instances",
+        "enable_metrics",
+        "query_warehouse",
     ],
 )
 def _submit_job(
@@ -447,6 +452,13 @@ def _submit_job(
         )
         target_instances = max(target_instances, kwargs.pop("num_instances"))
 
+    imports = None
+    if "additional_payloads" in kwargs:
+        logger.warning(
+            "'additional_payloads' is deprecated and will be removed in a future release. Use 'imports' instead."
+        )
+        imports = kwargs.pop("additional_payloads")
+
     # Use kwargs for less common optional parameters
     database = kwargs.pop("database", None)
     schema = kwargs.pop("schema", None)
@@ -457,10 +469,7 @@ def _submit_job(
     spec_overrides = kwargs.pop("spec_overrides", None)
     enable_metrics = kwargs.pop("enable_metrics", True)
     query_warehouse = kwargs.pop("query_warehouse", session.get_current_warehouse())
-    additional_payloads = kwargs.pop("additional_payloads", None)
-
-    if additional_payloads:
-        logger.warning("'additional_payloads' is in private preview since 1.9.1. Do not use it in production.")
+    imports = kwargs.pop("imports", None) or imports
 
     # Warn if there are unknown kwargs
     if kwargs:
@@ -492,7 +501,7 @@ def _submit_job(
     try:
         # Upload payload
         uploaded_payload = payload_utils.JobPayload(
-            source, entrypoint=entrypoint, pip_requirements=pip_requirements, additional_payloads=additional_payloads
+            source, entrypoint=entrypoint, pip_requirements=pip_requirements, additional_payloads=imports
         ).upload(session, stage_path)
     except snowpark.exceptions.SnowparkSQLException as e:
         if e.sql_error_code == 90106:
@@ -501,6 +510,22 @@ def _submit_job(
             )
         raise
 
+    # FIXME: Temporary patches, remove this after v1 is deprecated
+    if target_instances > 1:
+        default_spec_overrides = {
+            "spec": {
+                "endpoints": [
+                    {"name": "ray-dashboard-endpoint", "port": 12003, "protocol": "TCP"},
+                ]
+            },
+        }
+        if spec_overrides:
+            spec_overrides = spec_utils.merge_patch(
+                default_spec_overrides, spec_overrides, display_name="spec_overrides"
+            )
+        else:
+            spec_overrides = default_spec_overrides
+
     if feature_flags.FeatureFlags.USE_SUBMIT_JOB_V2.is_enabled():
         # Add default env vars (extracted from spec_utils.generate_service_spec)
         combined_env_vars = {**uploaded_payload.env_vars, **(env_vars or {})}
@@ -668,8 +693,10 @@ def _ensure_session(session: Optional[snowpark.Session]) -> snowpark.Session:
         session = session or get_active_session()
     except snowpark.exceptions.SnowparkSessionException as e:
         if "More than one active session" in e.message:
-            raise RuntimeError("Please specify the session as a parameter in API call")
+            raise RuntimeError(
+                "More than one active session is found. Please specify the session explicitly as a parameter"
+            ) from None
         if "No default Session is found" in e.message:
-            raise RuntimeError("Please create a session before API call")
+            raise RuntimeError("No active session is found. Please create a session") from None
         raise
     return session

snowflake/ml/lineage/lineage_node.py

@@ -83,7 +83,6 @@ class LineageNode(mixins.SerializableSessionMixin):
         raise NotImplementedError()
 
     @telemetry.send_api_usage_telemetry(project=_PROJECT)
-    @snowpark._internal.utils.private_preview(version="1.5.3")
     def lineage(
         self,
         direction: Literal["upstream", "downstream"] = "downstream",

snowflake/ml/model/__init__.py

@@ -1,10 +1,6 @@
-from snowflake.ml.model._client.model.batch_inference_specs import (
-    InputSpec,
-    JobSpec,
-    OutputSpec,
-)
+from snowflake.ml.model._client.model.batch_inference_specs import JobSpec, OutputSpec
 from snowflake.ml.model._client.model.model_impl import Model
 from snowflake.ml.model._client.model.model_version_impl import ExportMode, ModelVersion
 from snowflake.ml.model.models.huggingface_pipeline import HuggingFacePipelineModel
 
-__all__ = ["Model", "ModelVersion", "ExportMode", "HuggingFacePipelineModel", "InputSpec", "JobSpec", "OutputSpec"]
+__all__ = ["Model", "ModelVersion", "ExportMode", "HuggingFacePipelineModel", "JobSpec", "OutputSpec"]

snowflake/ml/model/_client/model/batch_inference_specs.py

@@ -3,10 +3,6 @@ from typing import Optional, Union
 from pydantic import BaseModel
 
 
-class InputSpec(BaseModel):
-    stage_location: str
-
-
 class OutputSpec(BaseModel):
     stage_location: str
 

snowflake/ml/model/_client/model/inference_engine_utils.py

@@ -0,0 +1,55 @@
+from typing import Any, Optional, Union
+
+from snowflake.ml.model._client.ops import service_ops
+
+
+def _get_inference_engine_args(
+    experimental_options: Optional[dict[str, Any]],
+) -> Optional[service_ops.InferenceEngineArgs]:
+
+    if not experimental_options:
+        return None
+
+    if "inference_engine" not in experimental_options:
+        raise ValueError("inference_engine is required in experimental_options")
+
+    return service_ops.InferenceEngineArgs(
+        inference_engine=experimental_options["inference_engine"],
+        inference_engine_args_override=experimental_options.get("inference_engine_args_override"),
+    )
+
+
+def _enrich_inference_engine_args(
+    inference_engine_args: service_ops.InferenceEngineArgs,
+    gpu_requests: Optional[Union[str, int]] = None,
+) -> Optional[service_ops.InferenceEngineArgs]:
+    """Enrich inference engine args with model path and tensor parallelism settings.
+
+    Args:
+        inference_engine_args: The original inference engine args
+        gpu_requests: The number of GPUs requested
+
+    Returns:
+        Enriched inference engine args
+
+    Raises:
+        ValueError: Invalid gpu_requests
+    """
+    if inference_engine_args.inference_engine_args_override is None:
+        inference_engine_args.inference_engine_args_override = []
+
+    gpu_count = None
+
+    # Set tensor-parallelism if gpu_requests is specified
+    if gpu_requests is not None:
+        # assert gpu_requests is a string or an integer before casting to int
+        try:
+            gpu_count = int(gpu_requests)
+            if gpu_count > 0:
+                inference_engine_args.inference_engine_args_override.append(f"--tensor-parallel-size={gpu_count}")
+            else:
+                raise ValueError(f"GPU count must be greater than 0, got {gpu_count}")
+        except ValueError:
+            raise ValueError(f"Invalid gpu_requests: {gpu_requests} with type {type(gpu_requests).__name__}")
+
+    return inference_engine_args