snowflake-ml-python 1.8.1__py3-none-any.whl → 1.8.3__py3-none-any.whl

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

@@ -0,0 +1,203 @@
+#!/usr/bin/env python3
+# This file is part of the Ray-based distributed job system for Snowflake ML.
+# Architecture overview:
+# - Head node creates a ShutdownSignal actor and signals workers when job completes
+# - Worker nodes listen for this signal and gracefully shut down
+# - This ensures clean termination of distributed Ray jobs
+import argparse
+import logging
+import socket
+import sys
+import time
+from typing import Any
+
+import ray
+from constants import (
+    SHUTDOWN_ACTOR_NAME,
+    SHUTDOWN_ACTOR_NAMESPACE,
+    SHUTDOWN_RPC_TIMEOUT_SECONDS,
+)
+from ray.actor import ActorHandle
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+
+
+@ray.remote
+class ShutdownSignal:
+    """A simple Ray actor that workers can check to determine if they should shutdown"""
+
+    def __init__(self) -> None:
+        self.shutdown_requested = False
+        self.timestamp = None
+        self.hostname = socket.gethostname()
+        self.acknowledged_workers = set()
+        logging.info(f"ShutdownSignal actor created on {self.hostname}")
+
+    def request_shutdown(self) -> dict[str, Any]:
+        """Signal workers to shut down"""
+        self.shutdown_requested = True
+        self.timestamp = time.time()
+        logging.info(f"Shutdown requested by head node at {self.timestamp}")
+        return {"status": "shutdown_requested", "timestamp": self.timestamp, "host": self.hostname}
+
+    def should_shutdown(self) -> dict[str, Any]:
+        """Check if shutdown has been requested"""
+        return {"shutdown": self.shutdown_requested, "timestamp": self.timestamp, "host": self.hostname}
+
+    def ping(self) -> dict[str, Any]:
+        """Simple method to test connectivity"""
+        return {"status": "alive", "host": self.hostname}
+
+    def acknowledge_shutdown(self, worker_id: str) -> dict[str, Any]:
+        """Worker acknowledges it has received the shutdown signal and is terminating"""
+        self.acknowledged_workers.add(worker_id)
+        logging.info(f"Worker {worker_id} acknowledged shutdown. Total acknowledged: {len(self.acknowledged_workers)}")
+
+        return {"status": "acknowledged", "worker_id": worker_id, "acknowledged_count": len(self.acknowledged_workers)}
+
+    def get_acknowledgment_workers(self) -> set[str]:
+        """Get the set of workers who have acknowledged shutdown"""
+        return self.acknowledged_workers
+
+
+def get_worker_node_ids() -> list[str]:
+    """Get the IDs of all active worker nodes.
+
+    Returns:
+        List[str]: List of worker node IDs. Empty list if no workers are present.
+    """
+    worker_nodes = [
+        node for node in ray.nodes() if node.get("Alive") and node.get("Resources", {}).get("node_tag:worker", 0) > 0
+    ]
+
+    worker_node_ids = [node.get("NodeName") for node in worker_nodes]
+
+    if worker_node_ids:
+        logging.info(f"Found {len(worker_node_ids)} worker nodes")
+    else:
+        logging.info("No active worker nodes found")
+
+    return worker_node_ids
+
+
+def get_or_create_shutdown_signal() -> ActorHandle:
+    """Get existing shutdown signal actor or create a new one.
+
+    Returns:
+        ActorHandle: Reference to shutdown signal actor
+    """
+    try:
+        # Try to get existing actor
+        shutdown_signal = ray.get_actor(SHUTDOWN_ACTOR_NAME, namespace=SHUTDOWN_ACTOR_NAMESPACE)
+        logging.info("Found existing shutdown signal actor")
+    except (ValueError, ray.exceptions.RayActorError) as e:
+        logging.info(f"Creating new shutdown signal actor: {e}")
+        # Create new actor if it doesn't exist
+        shutdown_signal = ShutdownSignal.options(
+            name=SHUTDOWN_ACTOR_NAME,
+            namespace=SHUTDOWN_ACTOR_NAMESPACE,
+            lifetime="detached",  # Ensure actor survives client disconnect
+            resources={"node_tag:head": 0.001},  # Resource constraint to ensure it runs on head node
+        ).remote()
+
+        # Verify actor is created and accessible
+        ping_result = ray.get(shutdown_signal.ping.remote(), timeout=SHUTDOWN_RPC_TIMEOUT_SECONDS)
+        logging.debug(f"New actor ping response: {ping_result}")
+
+    return shutdown_signal
+
+
+def request_shutdown(shutdown_signal: ActorHandle) -> None:
+    """Request workers to shut down.
+
+    Args:
+        shutdown_signal: Reference to the shutdown signal actor
+    """
+    response = ray.get(shutdown_signal.request_shutdown.remote(), timeout=SHUTDOWN_RPC_TIMEOUT_SECONDS)
+    logging.info(f"Shutdown requested: {response}")
+
+
+def verify_shutdown(shutdown_signal: ActorHandle) -> None:
+    """Verify that shutdown was properly signaled.
+
+    Args:
+        shutdown_signal: Reference to the shutdown signal actor
+    """
+    check = ray.get(shutdown_signal.should_shutdown.remote(), timeout=SHUTDOWN_RPC_TIMEOUT_SECONDS)
+    logging.debug(f"Shutdown status check: {check}")
+
+
+def wait_for_acknowledgments(shutdown_signal: ActorHandle, worker_node_ids: list[str], wait_time: int) -> None:
+    """Wait for workers to acknowledge shutdown.
+
+    Args:
+        shutdown_signal: Reference to the shutdown signal actor
+        worker_node_ids: List of worker node IDs
+        wait_time: Time in seconds to wait for acknowledgments
+
+    Raises:
+        TimeoutError: When workers don't acknowledge within the wait time or if actor communication times out
+    """
+    if not worker_node_ids:
+        return
+
+    logging.info(f"Waiting up to {wait_time}s for workers to acknowledge shutdown signal...")
+    start_time = time.time()
+    check_interval = 1.0
+
+    while time.time() - start_time < wait_time:
+        try:
+            ack_workers = ray.get(
+                shutdown_signal.get_acknowledgment_workers.remote(), timeout=SHUTDOWN_RPC_TIMEOUT_SECONDS
+            )
+            if ack_workers and ack_workers == set(worker_node_ids):
+                logging.info(
+                    f"All {len(worker_node_ids)} workers acknowledged shutdown. "
+                    f"Completed in {time.time() - start_time:.2f}s"
+                )
+                return
+            else:
+                logging.debug(f"Waiting for acknowledgments: {len(ack_workers)}/{len(worker_node_ids)} workers")
+        except Exception as e:
+            logging.warning(f"Error checking acknowledgment status: {e}")
+
+        time.sleep(check_interval)
+
+    raise TimeoutError(
+        f"Timed out waiting for {len(worker_node_ids)} workers to acknowledge shutdown after {wait_time}s"
+    )
+
+
+def signal_workers(wait_time: int = 10) -> int:
+    """
+    Signal worker nodes to shut down by creating a shutdown signal actor.
+
+    Args:
+        wait_time: Time in seconds to wait for workers to receive the message
+
+    Returns:
+        0 for success, 1 for failure
+    """
+    ray.init(address="auto", ignore_reinit_error=True)
+
+    worker_node_ids = get_worker_node_ids()
+
+    if worker_node_ids:
+        shutdown_signal = get_or_create_shutdown_signal()
+        request_shutdown(shutdown_signal)
+        verify_shutdown(shutdown_signal)
+        wait_for_acknowledgments(shutdown_signal, worker_node_ids, wait_time)
+    else:
+        logging.info("No active worker nodes found to signal.")
+
+    return 0
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Signal Ray workers to shutdown")
+    parser.add_argument(
+        "--wait-time", type=int, default=10, help="Time in seconds to wait for workers to receive the signal"
+    )
+    args = parser.parse_args()
+
+    sys.exit(signal_workers(args.wait_time))

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

@@ -0,0 +1,242 @@
+#!/usr/bin/env python3
+# This file is part of the Ray-based distributed job system for Snowflake ML.
+# Architecture overview:
+# - Head node creates a ShutdownSignal actor and signals workers when job completes
+# - Worker nodes listen for this signal via this script and gracefully shut down
+# - This ensures clean termination of distributed Ray jobs
+import logging
+import signal
+import sys
+import time
+from typing import Optional
+
+import get_instance_ip
+import ray
+from constants import (
+    SHUTDOWN_ACTOR_NAME,
+    SHUTDOWN_ACTOR_NAMESPACE,
+    SHUTDOWN_RPC_TIMEOUT_SECONDS,
+)
+from ray.actor import ActorHandle
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+
+
+def get_shutdown_actor() -> Optional[ActorHandle]:
+    """
+    Retrieve the shutdown signal actor from Ray.
+
+    Returns:
+        The shutdown signal actor or None if not found
+    """
+    try:
+        shutdown_signal = ray.get_actor(SHUTDOWN_ACTOR_NAME, namespace=SHUTDOWN_ACTOR_NAMESPACE)
+        return shutdown_signal
+    except Exception:
+        return None
+
+
+def ping_shutdown_actor(shutdown_signal: ActorHandle) -> bool:
+    """
+    Ping the shutdown actor to ensure connectivity.
+
+    Args:
+        shutdown_signal: The Ray actor handle for the shutdown signal
+
+    Returns:
+        True if ping succeeds, False otherwise
+    """
+    try:
+        ping_result = ray.get(shutdown_signal.ping.remote(), timeout=SHUTDOWN_RPC_TIMEOUT_SECONDS)
+        logging.debug(f"Actor ping result: {ping_result}")
+        return True
+    except (ray.exceptions.GetTimeoutError, Exception) as e:
+        logging.debug(f"Actor ping failed: {e}")
+        return False
+
+
+def check_shutdown_status(shutdown_signal: ActorHandle, worker_id: str) -> bool:
+    """
+    Check if worker should shutdown and acknowledge if needed.
+
+    Args:
+        shutdown_signal: The Ray actor handle for the shutdown signal
+        worker_id: Worker identifier (IP address)
+
+    Returns:
+        True if should shutdown, False otherwise
+    """
+    try:
+        status = ray.get(shutdown_signal.should_shutdown.remote(), timeout=SHUTDOWN_RPC_TIMEOUT_SECONDS)
+        logging.debug(f"Shutdown status: {status}")
+
+        if status.get("shutdown", False):
+            logging.info(
+                f"Received shutdown signal from head node at {status.get('timestamp')}. " f"Exiting worker process."
+            )
+
+            # Acknowledge shutdown before exiting
+            try:
+                ack_result = ray.get(
+                    shutdown_signal.acknowledge_shutdown.remote(worker_id), timeout=SHUTDOWN_RPC_TIMEOUT_SECONDS
+                )
+                logging.info(f"Acknowledged shutdown: {ack_result}")
+            except Exception as e:
+                logging.warning(f"Failed to acknowledge shutdown: {e}. Continue to exit worker.")
+
+            return True
+        return False
+
+    except Exception as e:
+        logging.debug(f"Error checking shutdown status: {e}")
+        return False
+
+
+def check_ray_connectivity() -> bool:
+    """
+    Check if the Ray cluster is accessible.
+
+    Returns:
+        True if Ray is connected, False otherwise
+    """
+    try:
+        # A simple check to verify Ray is working
+        nodes = ray.nodes()
+        if nodes:
+            return True
+        return False
+    except Exception as e:
+        logging.debug(f"Ray connectivity check failed: {e}")
+        return False
+
+
+def initialize_ray_connection(max_retries: int, initial_retry_delay: int, max_retry_delay: int) -> bool:
+    """
+    Initialize connection to Ray with retries.
+
+    Args:
+        max_retries: Maximum number of connection attempts
+        initial_retry_delay: Initial delay between retries in seconds
+        max_retry_delay: Maximum delay between retries in seconds
+
+    Returns:
+        bool: True if connection successful, False otherwise
+    """
+    retry_count = 0
+    retry_delay = initial_retry_delay
+
+    while retry_count < max_retries:
+        try:
+            ray.init(address="auto", ignore_reinit_error=True)
+            return True
+        except (ConnectionError, TimeoutError, RuntimeError) as e:
+            retry_count += 1
+            if retry_count >= max_retries:
+                logging.error(f"Failed to connect to Ray head after {max_retries} attempts: {e}")
+                return False
+
+            logging.debug(
+                f"Attempt {retry_count}/{max_retries} to connect to Ray failed: {e}. "
+                f"Retrying in {retry_delay} seconds..."
+            )
+            time.sleep(retry_delay)
+            # Exponential backoff with cap
+            retry_delay = min(retry_delay * 1.5, max_retry_delay)
+
+    return False  # Should not reach here, but added for completeness
+
+
+def monitor_shutdown_signal(check_interval: int, max_consecutive_failures: int) -> int:
+    """
+    Main loop to monitor for shutdown signals.
+
+    Args:
+        check_interval: Time in seconds between checks
+        max_consecutive_failures: Maximum allowed consecutive connection failures
+
+    Returns:
+        int: Exit code (0 for success, non-zero for failure)
+
+    Raises:
+        ConnectionError: If Ray connection failures exceed threshold
+    """
+    worker_id = get_instance_ip.get_self_ip()
+    actor_check_count = 0
+    consecutive_connection_failures = 0
+
+    logging.debug(
+        f"Starting to monitor for shutdown signal using actor {SHUTDOWN_ACTOR_NAME}"
+        f" in namespace {SHUTDOWN_ACTOR_NAMESPACE}."
+    )
+
+    while True:
+        actor_check_count += 1
+
+        # Check Ray connectivity before proceeding
+        if not check_ray_connectivity():
+            consecutive_connection_failures += 1
+            logging.debug(
+                f"Ray connectivity check failed (attempt {consecutive_connection_failures}/{max_consecutive_failures})"
+            )
+            if consecutive_connection_failures >= max_consecutive_failures:
+                raise ConnectionError("Exceeded max consecutive Ray connection failures")
+            time.sleep(check_interval)
+            continue
+
+        # Reset counter on successful connection
+        consecutive_connection_failures = 0
+
+        # Get shutdown actor
+        shutdown_signal = get_shutdown_actor()
+        if not shutdown_signal:
+            logging.debug(f"Shutdown signal actor not found at check #{actor_check_count}, continuing to wait...")
+            time.sleep(check_interval)
+            continue
+
+        # Ping the actor to ensure connectivity
+        if not ping_shutdown_actor(shutdown_signal):
+            time.sleep(check_interval)
+            continue
+
+        # Check shutdown status
+        if check_shutdown_status(shutdown_signal, worker_id):
+            return 0
+
+        # Wait before checking again
+        time.sleep(check_interval)
+
+
+def run_listener() -> int:
+    """Listen for shutdown signals from the head node"""
+    # Configuration
+    max_retries = 15
+    initial_retry_delay = 2
+    max_retry_delay = 30
+    check_interval = 5  # How often to check for ray connection or shutdown signal
+    max_consecutive_failures = 12  # Exit after about 1 minute of connection failures
+
+    # Initialize Ray connection
+    if not initialize_ray_connection(max_retries, initial_retry_delay, max_retry_delay):
+        raise ConnectionError("Failed to connect to Ray cluster. Aborting worker.")
+
+    # Monitor for shutdown signals
+    return monitor_shutdown_signal(check_interval, max_consecutive_failures)
+
+
+def main():
+    """Main entry point with signal handling"""
+
+    def signal_handler(signum, frame):
+        logging.info(f"Received signal {signum}, exiting worker process.")
+        sys.exit(0)
+
+    signal.signal(signal.SIGTERM, signal_handler)
+    signal.signal(signal.SIGINT, signal_handler)
+
+    # Run the listener - this will block until a shutdown signal is received
+    result = run_listener()
+    sys.exit(result)
+
+
+if __name__ == "__main__":
+    main()

snowflake/ml/jobs/_utils/spec_utils.py

@@ -1,7 +1,7 @@
 import logging
 from math import ceil
 from pathlib import PurePath
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Optional, Union
 
 from snowflake import snowpark
 from snowflake.ml._internal.utils import snowflake_env
@@ -15,10 +15,7 @@ def _get_node_resources(session: snowpark.Session, compute_pool: str) -> types.C
     if not rows:
         raise ValueError(f"Compute pool '{compute_pool}' not found")
     instance_family: str = rows[0]["instance_family"]
-
-    # Get the cloud we're using (AWS, Azure, etc)
-    region = snowflake_env.get_regions(session)[snowflake_env.get_current_region_id(session)]
-    cloud = region["cloud"]
+    cloud = snowflake_env.get_current_cloud(session, default=snowflake_env.SnowflakeCloudType.AWS)
 
     return (
         constants.COMMON_INSTANCE_FAMILIES.get(instance_family)
@@ -26,22 +23,14 @@ def _get_node_resources(session: snowpark.Session, compute_pool: str) -> types.C
     )
 
 
-def _get_image_spec(session: snowpark.Session, compute_pool: str, image_tag: Optional[str] = None) -> types.ImageSpec:
+def _get_image_spec(session: snowpark.Session, compute_pool: str) -> types.ImageSpec:
     # Retrieve compute pool node resources
     resources = _get_node_resources(session, compute_pool=compute_pool)
 
     # Use MLRuntime image
     image_repo = constants.DEFAULT_IMAGE_REPO
     image_name = constants.DEFAULT_IMAGE_GPU if resources.gpu > 0 else constants.DEFAULT_IMAGE_CPU
-
-    # Try to pull latest image tag from server side if possible
-    if not image_tag:
-        query_result = session.sql("SHOW PARAMETERS LIKE 'constants.RUNTIME_BASE_IMAGE_TAG' IN ACCOUNT").collect()
-        if query_result:
-            image_tag = query_result[0]["value"]
-
-    if image_tag is None:
-        image_tag = constants.DEFAULT_IMAGE_TAG
+    image_tag = constants.DEFAULT_IMAGE_TAG
 
     # TODO: Should each instance consume the entire pod?
     return types.ImageSpec(
@@ -54,9 +43,9 @@ def _get_image_spec(session: snowpark.Session, compute_pool: str, image_tag: Opt
 
 
 def generate_spec_overrides(
-    environment_vars: Optional[Dict[str, str]] = None,
-    custom_overrides: Optional[Dict[str, Any]] = None,
-) -> Dict[str, Any]:
+    environment_vars: Optional[dict[str, str]] = None,
+    custom_overrides: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
     """
     Generate a dictionary of service specification overrides.
 
@@ -68,7 +57,7 @@ def generate_spec_overrides(
         Resulting service specifiation patch dict. Empty if no overrides were supplied.
     """
     # Generate container level overrides
-    container_spec: Dict[str, Any] = {
+    container_spec: dict[str, Any] = {
         "name": constants.DEFAULT_CONTAINER_NAME,
     }
     if environment_vars:
@@ -95,10 +84,10 @@ def generate_service_spec(
     session: snowpark.Session,
     compute_pool: str,
     payload: types.UploadedPayload,
-    args: Optional[List[str]] = None,
+    args: Optional[list[str]] = None,
     num_instances: Optional[int] = None,
     enable_metrics: bool = False,
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     """
     Generate a service specification for a job.
 
@@ -114,20 +103,14 @@ def generate_service_spec(
         Job service specification
     """
     is_multi_node = num_instances is not None and num_instances > 1
+    image_spec = _get_image_spec(session, compute_pool)
 
     # Set resource requests/limits, including nvidia.com/gpu quantity if applicable
-    if is_multi_node:
-        # If the job is of multi-node, we will need a different image which contains
-        # module snowflake.runtime.utils.get_instance_ip
-        # TODO(SNOW-1961849): Remove the hard-coded image name
-        image_spec = _get_image_spec(session, compute_pool, constants.MULTINODE_HEADLESS_IMAGE_TAG)
-    else:
-        image_spec = _get_image_spec(session, compute_pool)
-    resource_requests: Dict[str, Union[str, int]] = {
+    resource_requests: dict[str, Union[str, int]] = {
         "cpu": f"{int(image_spec.resource_requests.cpu * 1000)}m",
         "memory": f"{image_spec.resource_limits.memory}Gi",
     }
-    resource_limits: Dict[str, Union[str, int]] = {
+    resource_limits: dict[str, Union[str, int]] = {
         "cpu": f"{int(image_spec.resource_requests.cpu * 1000)}m",
         "memory": f"{image_spec.resource_limits.memory}Gi",
     }
@@ -136,8 +119,8 @@ def generate_service_spec(
         resource_limits["nvidia.com/gpu"] = image_spec.resource_limits.gpu
 
     # Add local volumes for ephemeral logs and artifacts
-    volumes: List[Dict[str, str]] = []
-    volume_mounts: List[Dict[str, str]] = []
+    volumes: list[dict[str, str]] = []
+    volume_mounts: list[dict[str, str]] = []
     for volume_name, mount_path in [
         ("system-logs", "/var/log/managedservices/system/mlrs"),
         ("user-logs", "/var/log/managedservices/user/mlrs"),
@@ -191,7 +174,10 @@ def generate_service_spec(
 
     # TODO: Add hooks for endpoints for integration with TensorBoard etc
 
-    env_vars = {constants.PAYLOAD_DIR_ENV_VAR: stage_mount.as_posix()}
+    env_vars = {
+        constants.PAYLOAD_DIR_ENV_VAR: stage_mount.as_posix(),
+        constants.RESULT_PATH_ENV_VAR: constants.RESULT_PATH_DEFAULT_VALUE,
+    }
     endpoints = []
 
     if is_multi_node:
@@ -305,11 +291,11 @@ def merge_patch(base: Any, patch: Any, display_name: str = "") -> Any:
 
 
 def _merge_lists_of_dicts(
-    base: List[Dict[str, Any]],
-    patch: List[Dict[str, Any]],
+    base: list[dict[str, Any]],
+    patch: list[dict[str, Any]],
     merge_key: str = "name",
     display_name: str = "",
-) -> List[Dict[str, Any]]:
+) -> list[dict[str, Any]]:
     """
     Attempts to merge lists of dicts by matching on a merge key (default "name").
     - If the merge key is missing, the behavior falls back to overwriting the list.

snowflake/ml/jobs/_utils/types.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from pathlib import PurePath
-from typing import List, Literal, Optional, Union
+from typing import Literal, Optional, Union
 
 JOB_STATUS = Literal[
     "PENDING",
@@ -11,11 +11,17 @@ JOB_STATUS = Literal[
 ]
 
 
+@dataclass(frozen=True)
+class PayloadEntrypoint:
+    file_path: PurePath
+    main_func: Optional[str]
+
+
 @dataclass(frozen=True)
 class UploadedPayload:
     # TODO: Include manifest of payload files for validation
     stage_path: PurePath
-    entrypoint: List[Union[str, PurePath]]
+    entrypoint: list[Union[str, PurePath]]
 
 
 @dataclass(frozen=True)

snowflake/ml/jobs/decorators.py

@@ -1,6 +1,6 @@
 import copy
 import functools
-from typing import Callable, Dict, List, Optional, TypeVar
+from typing import Callable, Optional, TypeVar
 
 from typing_extensions import ParamSpec
 
@@ -15,20 +15,19 @@ _Args = ParamSpec("_Args")
 _ReturnValue = TypeVar("_ReturnValue")
 
 
-@snowpark._internal.utils.private_preview(version="1.7.4")
 @telemetry.send_api_usage_telemetry(project=_PROJECT)
 def remote(
     compute_pool: str,
     *,
     stage_name: str,
-    pip_requirements: Optional[List[str]] = None,
-    external_access_integrations: Optional[List[str]] = None,
+    pip_requirements: Optional[list[str]] = None,
+    external_access_integrations: Optional[list[str]] = None,
     query_warehouse: Optional[str] = None,
-    env_vars: Optional[Dict[str, str]] = None,
+    env_vars: Optional[dict[str, str]] = None,
     num_instances: Optional[int] = None,
     enable_metrics: bool = False,
     session: Optional[snowpark.Session] = None,
-) -> Callable[[Callable[_Args, _ReturnValue]], Callable[_Args, jb.MLJob]]:
+) -> Callable[[Callable[_Args, _ReturnValue]], Callable[_Args, jb.MLJob[_ReturnValue]]]:
     """
     Submit a job to the compute pool.
 
@@ -47,7 +46,7 @@ def remote(
         Decorator that dispatches invocations of the decorated function as remote jobs.
     """
 
-    def decorator(func: Callable[_Args, _ReturnValue]) -> Callable[_Args, jb.MLJob]:
+    def decorator(func: Callable[_Args, _ReturnValue]) -> Callable[_Args, jb.MLJob[_ReturnValue]]:
         # Copy the function to avoid modifying the original
         # We need to modify the line number of the function to exclude the
         # decorator from the copied source code
@@ -55,7 +54,7 @@ def remote(
         wrapped_func.__code__ = wrapped_func.__code__.replace(co_firstlineno=func.__code__.co_firstlineno + 1)
 
         @functools.wraps(func)
-        def wrapper(*args: _Args.args, **kwargs: _Args.kwargs) -> jb.MLJob:
+        def wrapper(*args: _Args.args, **kwargs: _Args.kwargs) -> jb.MLJob[_ReturnValue]:
             payload = functools.partial(func, *args, **kwargs)
             setattr(payload, constants.IS_MLJOB_REMOTE_ATTR, True)
             job = jm._submit_job(