snowflake-ml-python 1.11.0__py3-none-any.whl → 1.12.0__py3-none-any.whl

snowflake/ml/jobs/manager.py

@@ -1,6 +1,8 @@
+import json
 import logging
 import pathlib
 import textwrap
+from pathlib import PurePath
 from typing import Any, Callable, Optional, TypeVar, Union, cast, overload
 from uuid import uuid4
 
@@ -11,7 +13,13 @@ from snowflake import snowpark
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.utils import identifier
 from snowflake.ml.jobs import job as jb
-from snowflake.ml.jobs._utils import payload_utils, query_helper, spec_utils
+from snowflake.ml.jobs._utils import (
+    feature_flags,
+    payload_utils,
+    query_helper,
+    spec_utils,
+    types,
+)
 from snowflake.snowpark.context import get_active_session
 from snowflake.snowpark.exceptions import SnowparkSQLException
 from snowflake.snowpark.functions import coalesce, col, lit, when
@@ -445,7 +453,7 @@ def _submit_job(
     env_vars = kwargs.pop("env_vars", None)
     spec_overrides = kwargs.pop("spec_overrides", None)
     enable_metrics = kwargs.pop("enable_metrics", True)
-    query_warehouse = kwargs.pop("query_warehouse", None)
+    query_warehouse = kwargs.pop("query_warehouse", session.get_current_warehouse())
     additional_payloads = kwargs.pop("additional_payloads", None)
 
     if additional_payloads:
@@ -483,6 +491,27 @@ def _submit_job(
         source, entrypoint=entrypoint, pip_requirements=pip_requirements, additional_payloads=additional_payloads
     ).upload(session, stage_path)
 
+    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 {})}
+
+        return _do_submit_job_v2(
+            session=session,
+            payload=uploaded_payload,
+            args=args,
+            env_vars=combined_env_vars,
+            spec_overrides=spec_overrides,
+            compute_pool=compute_pool,
+            job_id=job_id,
+            external_access_integrations=external_access_integrations,
+            query_warehouse=query_warehouse,
+            target_instances=target_instances,
+            min_instances=min_instances,
+            enable_metrics=enable_metrics,
+            use_async=True,
+        )
+
+    # Fall back to v1
     # Generate service spec
     spec = spec_utils.generate_service_spec(
         session,
@@ -493,6 +522,8 @@ def _submit_job(
         min_instances=min_instances,
         enable_metrics=enable_metrics,
     )
+
+    # Generate spec overrides
     spec_overrides = spec_utils.generate_spec_overrides(
         environment_vars=env_vars,
         custom_overrides=spec_overrides,
@@ -500,26 +531,25 @@ def _submit_job(
     if spec_overrides:
         spec = spec_utils.merge_patch(spec, spec_overrides, display_name="spec_overrides")
 
-    query_text, params = _generate_submission_query(
-        spec, external_access_integrations, query_warehouse, target_instances, session, compute_pool, job_id
+    return _do_submit_job_v1(
+        session, spec, external_access_integrations, query_warehouse, target_instances, compute_pool, job_id
     )
-    _ = query_helper.run_query(session, query_text, params=params)
-    return get_job(job_id, session=session)
 
 
-def _generate_submission_query(
+def _do_submit_job_v1(
+    session: snowpark.Session,
     spec: dict[str, Any],
     external_access_integrations: list[str],
     query_warehouse: Optional[str],
     target_instances: int,
-    session: snowpark.Session,
     compute_pool: str,
     job_id: str,
-) -> tuple[str, list[Any]]:
+) -> jb.MLJob[Any]:
     """
     Generate the SQL query for job submission.
 
     Args:
+        session: The Snowpark session to use.
         spec: The service spec for the job.
         external_access_integrations: The external access integrations for the job.
         query_warehouse: The query warehouse for the job.
@@ -529,7 +559,7 @@ def _generate_submission_query(
         job_id: The ID of the job.
 
     Returns:
-        A tuple containing the SQL query text and the parameters for the query.
+        The job object.
     """
     query_template = textwrap.dedent(
         """\
@@ -547,12 +577,77 @@ def _generate_submission_query(
     if external_access_integrations:
         external_access_integration_list = ",".join(f"{e}" for e in external_access_integrations)
         query.append(f"EXTERNAL_ACCESS_INTEGRATIONS = ({external_access_integration_list})")
-    query_warehouse = query_warehouse or session.get_current_warehouse()
     if query_warehouse:
         query.append("QUERY_WAREHOUSE = IDENTIFIER(?)")
         params.append(query_warehouse)
     if target_instances > 1:
         query.append("REPLICAS = ?")
         params.append(target_instances)
+
     query_text = "\n".join(line for line in query if line)
-    return query_text, params
+    _ = query_helper.run_query(session, query_text, params=params)
+
+    return get_job(job_id, session=session)
+
+
+def _do_submit_job_v2(
+    session: snowpark.Session,
+    payload: types.UploadedPayload,
+    args: Optional[list[str]],
+    env_vars: dict[str, str],
+    spec_overrides: dict[str, Any],
+    compute_pool: str,
+    job_id: Optional[str] = None,
+    external_access_integrations: Optional[list[str]] = None,
+    query_warehouse: Optional[str] = None,
+    target_instances: int = 1,
+    min_instances: int = 1,
+    enable_metrics: bool = True,
+    use_async: bool = True,
+) -> jb.MLJob[Any]:
+    """
+    Generate the SQL query for job submission.
+
+    Args:
+        session: The Snowpark session to use.
+        payload: The uploaded job payload.
+        args: Arguments to pass to the entrypoint script.
+        env_vars: Environment variables to set in the job container.
+        spec_overrides: Custom service specification overrides.
+        compute_pool: The compute pool to use for job execution.
+        job_id: The ID of the job.
+        external_access_integrations: Optional list of external access integrations.
+        query_warehouse: Optional query warehouse to use.
+        target_instances: Number of instances for multi-node job.
+        min_instances: Minimum number of instances required to start the job.
+        enable_metrics: Whether to enable platform metrics for the job.
+        use_async: Whether to run the job asynchronously.
+
+    Returns:
+        The job object.
+    """
+    args = [
+        (payload.stage_path.joinpath(v).as_posix() if isinstance(v, PurePath) else v) for v in payload.entrypoint
+    ] + (args or [])
+    spec_options = {
+        "STAGE_PATH": payload.stage_path.as_posix(),
+        "ENTRYPOINT": ["/usr/local/bin/_entrypoint.sh"],
+        "ARGS": args,
+        "ENV_VARS": env_vars,
+        "ENABLE_METRICS": enable_metrics,
+        "SPEC_OVERRIDES": spec_overrides,
+    }
+    job_options = {
+        "EXTERNAL_ACCESS_INTEGRATIONS": external_access_integrations,
+        "QUERY_WAREHOUSE": query_warehouse,
+        "TARGET_INSTANCES": target_instances,
+        "MIN_INSTANCES": min_instances,
+        "ASYNC": use_async,
+    }
+    job_options = {k: v for k, v in job_options.items() if v is not None}
+
+    query_template = "CALL SYSTEM$EXECUTE_ML_JOB(?, ?, ?, ?)"
+    params = [job_id, compute_pool, json.dumps(spec_options), json.dumps(job_options)]
+    actual_job_id = query_helper.run_query(session, query_template, params=params)[0][0]
+
+    return get_job(actual_job_id, session=session)

snowflake/ml/model/__init__.py

@@ -1,5 +1,10 @@
+from snowflake.ml.model._client.model.batch_inference_specs import (
+    InputSpec,
+    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"]
+__all__ = ["Model", "ModelVersion", "ExportMode", "HuggingFacePipelineModel", "InputSpec", "JobSpec", "OutputSpec"]

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

@@ -0,0 +1,27 @@
+from typing import Optional, Union
+
+from pydantic import BaseModel
+
+
+class InputSpec(BaseModel):
+    input_stage_location: str
+    input_file_pattern: str = "*"
+
+
+class OutputSpec(BaseModel):
+    output_stage_location: str
+    output_file_prefix: Optional[str] = None
+    completion_filename: str = "_SUCCESS"
+
+
+class JobSpec(BaseModel):
+    image_repo: Optional[str] = None
+    job_name: Optional[str] = None
+    num_workers: Optional[int] = None
+    function_name: Optional[str] = None
+    gpu: Optional[Union[str, int]] = None
+    force_rebuild: bool = False
+    max_batch_rows: int = 1024
+    warehouse: Optional[str] = None
+    cpu_requests: Optional[str] = None
+    memory_requests: Optional[str] = None

snowflake/ml/model/_client/model/model_version_impl.py

@@ -1,16 +1,18 @@
 import enum
 import pathlib
 import tempfile
+import uuid
 import warnings
 from typing import Any, Callable, Optional, Union, overload
 
 import pandas as pd
 
-from snowflake import snowpark
+from snowflake.ml import jobs
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.utils import sql_identifier
 from snowflake.ml.lineage import lineage_node
 from snowflake.ml.model import task, type_hints
+from snowflake.ml.model._client.model import batch_inference_specs
 from snowflake.ml.model._client.ops import metadata_ops, model_ops, service_ops
 from snowflake.ml.model._model_composer import model_composer
 from snowflake.ml.model._model_composer.model_manifest import model_manifest_schema
@@ -19,6 +21,7 @@ from snowflake.snowpark import Session, async_job, dataframe
 
 _TELEMETRY_PROJECT = "MLOps"
 _TELEMETRY_SUBPROJECT = "ModelManagement"
+_BATCH_INFERENCE_JOB_ID_PREFIX = "BATCH_INFERENCE_"
 
 
 class ExportMode(enum.Enum):
@@ -539,6 +542,63 @@ class ModelVersion(lineage_node.LineageNode):
                 is_partitioned=target_function_info["is_partitioned"],
             )
 
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+        func_params_to_log=[
+            "compute_pool",
+        ],
+    )
+    def _run_batch(
+        self,
+        *,
+        compute_pool: str,
+        input_spec: batch_inference_specs.InputSpec,
+        output_spec: batch_inference_specs.OutputSpec,
+        job_spec: Optional[batch_inference_specs.JobSpec] = None,
+    ) -> jobs.MLJob[Any]:
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+
+        if job_spec is None:
+            job_spec = batch_inference_specs.JobSpec()
+
+        warehouse = job_spec.warehouse or self._service_ops._session.get_current_warehouse()
+        if warehouse is None:
+            raise ValueError("Warehouse is not set. Please set the warehouse field in the JobSpec.")
+
+        if job_spec.job_name is None:
+            # Same as the MLJob ID generation logic with a different prefix
+            job_name = f"{_BATCH_INFERENCE_JOB_ID_PREFIX}{str(uuid.uuid4()).replace('-', '_').upper()}"
+        else:
+            job_name = job_spec.job_name
+
+        return self._service_ops.invoke_batch_job_method(
+            # model version info
+            model_name=self._model_name,
+            version_name=self._version_name,
+            # job spec
+            function_name=self._get_function_info(function_name=job_spec.function_name)["target_method"],
+            compute_pool_name=sql_identifier.SqlIdentifier(compute_pool),
+            force_rebuild=job_spec.force_rebuild,
+            image_repo_name=job_spec.image_repo,
+            num_workers=job_spec.num_workers,
+            max_batch_rows=job_spec.max_batch_rows,
+            warehouse=sql_identifier.SqlIdentifier(warehouse),
+            cpu_requests=job_spec.cpu_requests,
+            memory_requests=job_spec.memory_requests,
+            job_name=job_name,
+            # input and output
+            input_stage_location=input_spec.input_stage_location,
+            input_file_pattern=input_spec.input_file_pattern,
+            output_stage_location=output_spec.output_stage_location,
+            completion_filename=output_spec.completion_filename,
+            # misc
+            statement_params=statement_params,
+        )
+
     def _get_function_info(self, function_name: Optional[str]) -> model_manifest_schema.ModelFunctionInfo:
         functions: list[model_manifest_schema.ModelFunctionInfo] = self._functions
 
@@ -1184,69 +1244,5 @@ class ModelVersion(lineage_node.LineageNode):
             statement_params=statement_params,
         )
 
-    @snowpark._internal.utils.private_preview(version="1.8.3")
-    @telemetry.send_api_usage_telemetry(
-        project=_TELEMETRY_PROJECT,
-        subproject=_TELEMETRY_SUBPROJECT,
-    )
-    def _run_job(
-        self,
-        X: Union[pd.DataFrame, "dataframe.DataFrame"],
-        *,
-        job_name: str,
-        compute_pool: str,
-        image_repo: Optional[str] = None,
-        output_table_name: str,
-        function_name: Optional[str] = None,
-        cpu_requests: Optional[str] = None,
-        memory_requests: Optional[str] = None,
-        gpu_requests: Optional[Union[str, int]] = None,
-        num_workers: Optional[int] = None,
-        max_batch_rows: Optional[int] = None,
-        force_rebuild: bool = False,
-        build_external_access_integrations: Optional[list[str]] = None,
-    ) -> Union[pd.DataFrame, dataframe.DataFrame]:
-        statement_params = telemetry.get_statement_params(
-            project=_TELEMETRY_PROJECT,
-            subproject=_TELEMETRY_SUBPROJECT,
-        )
-        target_function_info = self._get_function_info(function_name=function_name)
-        job_db_id, job_schema_id, job_id = sql_identifier.parse_fully_qualified_name(job_name)
-        output_table_db_id, output_table_schema_id, output_table_id = sql_identifier.parse_fully_qualified_name(
-            output_table_name
-        )
-        warehouse = self._service_ops._session.get_current_warehouse()
-        assert warehouse, "No active warehouse selected in the current session."
-        return self._service_ops.invoke_job_method(
-            target_method=target_function_info["target_method"],
-            signature=target_function_info["signature"],
-            X=X,
-            database_name=None,
-            schema_name=None,
-            model_name=self._model_name,
-            version_name=self._version_name,
-            job_database_name=job_db_id,
-            job_schema_name=job_schema_id,
-            job_name=job_id,
-            compute_pool_name=sql_identifier.SqlIdentifier(compute_pool),
-            warehouse_name=sql_identifier.SqlIdentifier(warehouse),
-            image_repo_name=image_repo,
-            output_table_database_name=output_table_db_id,
-            output_table_schema_name=output_table_schema_id,
-            output_table_name=output_table_id,
-            cpu_requests=cpu_requests,
-            memory_requests=memory_requests,
-            gpu_requests=gpu_requests,
-            num_workers=num_workers,
-            max_batch_rows=max_batch_rows,
-            force_rebuild=force_rebuild,
-            build_external_access_integrations=(
-                None
-                if build_external_access_integrations is None
-                else [sql_identifier.SqlIdentifier(eai) for eai in build_external_access_integrations]
-            ),
-            statement_params=statement_params,
-        )
-
 
 lineage_node.DOMAIN_LINEAGE_REGISTRY["model"] = ModelVersion

snowflake/ml/model/_client/ops/service_ops.py

@@ -10,17 +10,13 @@ import time
 from typing import Any, Optional, Union, cast
 
 from snowflake import snowpark
+from snowflake.ml import jobs
 from snowflake.ml._internal import file_utils, platform_capabilities as pc
 from snowflake.ml._internal.utils import identifier, service_logger, sql_identifier
-from snowflake.ml.model import (
-    inference_engine as inference_engine_module,
-    model_signature,
-    type_hints,
-)
+from snowflake.ml.model import inference_engine as inference_engine_module, type_hints
 from snowflake.ml.model._client.service import model_deployment_spec
 from snowflake.ml.model._client.sql import service as service_sql, stage as stage_sql
-from snowflake.ml.model._signatures import snowpark_handler
-from snowflake.snowpark import async_job, dataframe, exceptions, row, session
+from snowflake.snowpark import async_job, exceptions, row, session
 from snowflake.snowpark._internal import utils as snowpark_utils
 
 module_logger = service_logger.get_logger(__name__, service_logger.LogColor.GREY)
@@ -866,174 +862,97 @@ class ServiceOperator:
         except exceptions.SnowparkSQLException:
             return False
 
-    def invoke_job_method(
+    def invoke_batch_job_method(
         self,
-        target_method: str,
-        signature: model_signature.ModelSignature,
-        X: Union[type_hints.SupportedDataType, dataframe.DataFrame],
-        database_name: Optional[sql_identifier.SqlIdentifier],
-        schema_name: Optional[sql_identifier.SqlIdentifier],
+        *,
+        function_name: str,
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        job_database_name: Optional[sql_identifier.SqlIdentifier],
-        job_schema_name: Optional[sql_identifier.SqlIdentifier],
-        job_name: sql_identifier.SqlIdentifier,
+        job_name: str,
         compute_pool_name: sql_identifier.SqlIdentifier,
-        warehouse_name: sql_identifier.SqlIdentifier,
+        warehouse: sql_identifier.SqlIdentifier,
         image_repo_name: Optional[str],
-        output_table_database_name: Optional[sql_identifier.SqlIdentifier],
-        output_table_schema_name: Optional[sql_identifier.SqlIdentifier],
-        output_table_name: sql_identifier.SqlIdentifier,
-        cpu_requests: Optional[str],
-        memory_requests: Optional[str],
-        gpu_requests: Optional[Union[int, str]],
+        input_stage_location: str,
+        input_file_pattern: str,
+        output_stage_location: str,
+        completion_filename: str,
+        force_rebuild: bool,
         num_workers: Optional[int],
         max_batch_rows: Optional[int],
-        force_rebuild: bool,
-        build_external_access_integrations: Optional[list[sql_identifier.SqlIdentifier]],
+        cpu_requests: Optional[str],
+        memory_requests: Optional[str],
         statement_params: Optional[dict[str, Any]] = None,
-    ) -> Union[type_hints.SupportedDataType, dataframe.DataFrame]:
-        # fall back to the registry's database and schema if not provided
-        database_name = database_name or self._database_name
-        schema_name = schema_name or self._schema_name
-
-        # fall back to the model's database and schema if not provided then to the registry's database and schema
-        job_database_name = job_database_name or database_name or self._database_name
-        job_schema_name = job_schema_name or schema_name or self._schema_name
+    ) -> jobs.MLJob[Any]:
+        database_name = self._database_name
+        schema_name = self._schema_name
 
-        image_repo_fqn = self._get_image_repo_fqn(image_repo_name, database_name, schema_name)
+        job_database_name, job_schema_name, job_name = sql_identifier.parse_fully_qualified_name(job_name)
+        job_database_name = job_database_name or database_name
+        job_schema_name = job_schema_name or schema_name
 
-        input_table_database_name = job_database_name
-        input_table_schema_name = job_schema_name
-        output_table_database_name = output_table_database_name or database_name or self._database_name
-        output_table_schema_name = output_table_schema_name or schema_name or self._schema_name
-
-        if self._workspace:
-            stage_path = self._create_temp_stage(database_name, schema_name, statement_params)
-        else:
-            stage_path = None
+        self._model_deployment_spec.clear()
 
-        # validate and prepare input
-        if not isinstance(X, dataframe.DataFrame):
-            keep_order = True
-            output_with_input_features = False
-            df = model_signature._convert_and_validate_local_data(X, signature.inputs)
-            s_df = snowpark_handler.SnowparkDataFrameHandler.convert_from_df(
-                self._session, df, keep_order=keep_order, features=signature.inputs, statement_params=statement_params
-            )
-        else:
-            keep_order = False
-            output_with_input_features = True
-            s_df = X
-
-        # only write the index and feature input columns
-        cols = [snowpark_handler._KEEP_ORDER_COL_NAME] if snowpark_handler._KEEP_ORDER_COL_NAME in s_df.columns else []
-        cols += [
-            sql_identifier.SqlIdentifier(feature.name, case_sensitive=True).identifier() for feature in signature.inputs
-        ]
-        s_df = s_df.select(cols)
-        original_cols = s_df.columns
-
-        # input/output tables
-        fq_output_table_name = identifier.get_schema_level_object_identifier(
-            output_table_database_name.identifier(),
-            output_table_schema_name.identifier(),
-            output_table_name.identifier(),
-        )
-        tmp_input_table_id = sql_identifier.SqlIdentifier(
-            snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.TABLE)
-        )
-        fq_tmp_input_table_name = identifier.get_schema_level_object_identifier(
-            job_database_name.identifier(),
-            job_schema_name.identifier(),
-            tmp_input_table_id.identifier(),
-        )
-        s_df.write.save_as_table(
-            table_name=fq_tmp_input_table_name,
-            mode="errorifexists",
-            statement_params=statement_params,
+        self._model_deployment_spec.add_model_spec(
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            version_name=version_name,
         )
 
-        try:
-            self._model_deployment_spec.clear()
-            # save the spec
-            self._model_deployment_spec.add_model_spec(
-                database_name=database_name,
-                schema_name=schema_name,
-                model_name=model_name,
-                version_name=version_name,
-            )
-            self._model_deployment_spec.add_job_spec(
-                job_database_name=job_database_name,
-                job_schema_name=job_schema_name,
-                job_name=job_name,
-                inference_compute_pool_name=compute_pool_name,
-                cpu=cpu_requests,
-                memory=memory_requests,
-                gpu=gpu_requests,
-                num_workers=num_workers,
-                max_batch_rows=max_batch_rows,
-                warehouse=warehouse_name,
-                target_method=target_method,
-                input_table_database_name=input_table_database_name,
-                input_table_schema_name=input_table_schema_name,
-                input_table_name=tmp_input_table_id,
-                output_table_database_name=output_table_database_name,
-                output_table_schema_name=output_table_schema_name,
-                output_table_name=output_table_name,
-            )
+        self._model_deployment_spec.add_job_spec(
+            job_database_name=job_database_name,
+            job_schema_name=job_schema_name,
+            job_name=job_name,
+            inference_compute_pool_name=compute_pool_name,
+            num_workers=num_workers,
+            max_batch_rows=max_batch_rows,
+            input_stage_location=input_stage_location,
+            input_file_pattern=input_file_pattern,
+            output_stage_location=output_stage_location,
+            completion_filename=completion_filename,
+            function_name=function_name,
+            warehouse=warehouse,
+            cpu=cpu_requests,
+            memory=memory_requests,
+        )
 
-            self._model_deployment_spec.add_image_build_spec(
-                image_build_compute_pool_name=compute_pool_name,
-                fully_qualified_image_repo_name=image_repo_fqn,
-                force_rebuild=force_rebuild,
-                external_access_integrations=build_external_access_integrations,
-            )
+        self._model_deployment_spec.add_image_build_spec(
+            image_build_compute_pool_name=compute_pool_name,
+            fully_qualified_image_repo_name=self._get_image_repo_fqn(image_repo_name, database_name, schema_name),
+            force_rebuild=force_rebuild,
+        )
 
-            spec_yaml_str_or_path = self._model_deployment_spec.save()
-            if self._workspace:
-                assert stage_path is not None
-                file_utils.upload_directory_to_stage(
-                    self._session,
-                    local_path=pathlib.Path(self._workspace.name),
-                    stage_path=pathlib.PurePosixPath(stage_path),
-                    statement_params=statement_params,
-                )
+        spec_yaml_str_or_path = self._model_deployment_spec.save()
 
-            # deploy the job
-            query_id, async_job = self._service_client.deploy_model(
-                stage_path=stage_path if self._workspace else None,
-                model_deployment_spec_file_rel_path=(
-                    model_deployment_spec.ModelDeploymentSpec.DEPLOY_SPEC_FILE_REL_PATH if self._workspace else None
-                ),
-                model_deployment_spec_yaml_str=None if self._workspace else spec_yaml_str_or_path,
+        if self._workspace:
+            module_logger.info("using workspace")
+            stage_path = self._create_temp_stage(database_name, schema_name, statement_params)
+            file_utils.upload_directory_to_stage(
+                self._session,
+                local_path=pathlib.Path(self._workspace.name),
+                stage_path=pathlib.PurePosixPath(stage_path),
                 statement_params=statement_params,
             )
+        else:
+            module_logger.info("not using workspace")
+            stage_path = None
 
-            while not async_job.is_done():
-                time.sleep(5)
-        finally:
-            self._session.table(fq_tmp_input_table_name).drop_table()
-
-        # handle the output
-        df_res = self._session.table(fq_output_table_name)
-        if keep_order:
-            df_res = df_res.sort(
-                snowpark_handler._KEEP_ORDER_COL_NAME,
-                ascending=True,
-            )
-            df_res = df_res.drop(snowpark_handler._KEEP_ORDER_COL_NAME)
+        _, async_job = self._service_client.deploy_model(
+            stage_path=stage_path if self._workspace else None,
+            model_deployment_spec_file_rel_path=(
+                model_deployment_spec.ModelDeploymentSpec.DEPLOY_SPEC_FILE_REL_PATH if self._workspace else None
+            ),
+            model_deployment_spec_yaml_str=None if self._workspace else spec_yaml_str_or_path,
+            statement_params=statement_params,
+        )
 
-        if not output_with_input_features:
-            df_res = df_res.drop(*original_cols)
+        # Block until the async job is done
+        async_job.result()
 
-        # get final result
-        if not isinstance(X, dataframe.DataFrame):
-            return snowpark_handler.SnowparkDataFrameHandler.convert_to_df(
-                df_res, features=signature.outputs, statement_params=statement_params
-            )
-        else:
-            return df_res
+        return jobs.MLJob(
+            id=sql_identifier.get_fully_qualified_name(job_database_name, job_schema_name, job_name),
+            session=self._session,
+        )
 
     def _create_temp_stage(
         self,