fix: enable SLS_CORE with envvar RUNPOD_SLS_CORE=1

Author: ef0xa

runpod/serverless/__init__.py

@@ -4,21 +4,22 @@
 Arguments can be passed in when the worker is started, and will be passed to the worker.
 """
 
+import argparse
+import json
 import os
+import signal
 import sys
-import json
 import time
-import signal
-import argparse
-from typing import Dict, Any
+import typing
+from typing import Any, Dict
 
 from runpod.serverless import core
+
+from ..version import __version__ as runpod_version
 from . import worker
 from .modules import rp_fastapi
 from .modules.rp_logger import RunPodLogger
 from .modules.rp_progress import progress_update
-from ..version import __version__ as runpod_version
-
 
 log = RunPodLogger()
 
@@ -65,11 +66,13 @@ def _set_config_args(config) -> dict:
 
     # Parse the test input from JSON
     if config["rp_args"]["test_input"]:
-        config["rp_args"]["test_input"] = json.loads(config["rp_args"]["test_input"])
+        config["rp_args"]["test_input"] = json.loads(
+            config["rp_args"]["test_input"])
 
     # Parse the test output from JSON
     if config["rp_args"].get("test_output", None):
-        config["rp_args"]["test_output"] = json.loads(config["rp_args"]["test_output"])
+        config["rp_args"]["test_output"] = json.loads(
+            config["rp_args"]["test_output"])
 
     # Set the log level
     if config["rp_args"]["rp_log_level"]:
@@ -133,8 +136,9 @@ def start(config: Dict[str, Any]):
             api_port=config['rp_args']['rp_api_port'],
             api_concurrency=config['rp_args']['rp_api_concurrency']
         )
+        return
 
-    elif realtime_port:
+    if realtime_port:
         log.info(f"Starting API server for realtime on port {realtime_port}.")
         api_server = rp_fastapi.WorkerAPI(config)
 
@@ -143,12 +147,13 @@ def start(config: Dict[str, Any]):
             api_port=realtime_port,
             api_concurrency=realtime_concurrency
         )
+        return
 
     # --------------------------------- SLS-Core --------------------------------- #
-    elif os.environ.get("RUNPOD_USE_CORE", None) or os.environ.get("RUNPOD_CORE_PATH", None):
-        log.info("Starting worker with SLS-Core.")
+    if os.getenv("RUNPOD_SLS_CORE", "false").lower() in ("1", 't', 'T', 'TRUE', 'true', 'True', '0', 'f', 'F', 'FALSE', 'false', 'False'):
         core.main(config)
+        return
 
     # --------------------------------- Standard --------------------------------- #
-    else:
-        worker.main(config)
+    worker.main(config)
+    return

runpod/serverless/core.py

@@ -8,6 +8,7 @@
 import asyncio
 from ctypes import CDLL, byref, c_char_p, c_int
 from typing import Any, Callable,  List, Dict, Optional
+import typing
 
 from runpod.version import __version__ as runpod_version
 from runpod.serverless.modules.rp_logger import RunPodLogger
@@ -31,28 +32,35 @@ def __str__(self) -> str:
         return f"CGetJobResult(res_len={self.res_len}, status_code={self.status_code})"
 
 
+def notregistered():
+    """ Function to raise NotImplementedError """
+    raise RuntimeError("This function is not registered with the SLS Core.")
 class Hook:  # pylint: disable=too-many-instance-attributes
     """ Singleton class for interacting with sls_core.so"""
 
     _instance = None
 
+
     # C function pointers
-    _get_jobs: Callable = None
-    _progress_update: Callable = None
-    _stream_output: Callable = None
-    _post_output: Callable = None
-    _finish_stream: Callable = None
+    _get_jobs: Callable = notregistered
+    _progress_update: Callable = notregistered
+    _stream_output: Callable = notregistered
+    _post_output: Callable = notregistered
+    _finish_stream: Callable = notregistered
 
     def __new__(cls):
         if Hook._instance is None:
             log.debug("SLS Core | Initializing Hook.")
             Hook._instance = object.__new__(cls)
             Hook._initialized = False
+        
         return Hook._instance
 
     def __init__(self, rust_so_path: Optional[str] = None) -> None:
+        
         if self._initialized:
             return
+        
 
         if rust_so_path is None:
             default_path = os.path.join(
@@ -169,8 +177,6 @@ def finish_stream(self, job_id: str) -> bool:
         return bool(self._finish_stream(
             c_char_p(id_bytes), c_int(len(id_bytes))
         ))
-
-
 # -------------------------------- Process Job ------------------------------- #
 async def _process_job(config: Dict[str, Any], job: Dict[str, Any], hook) -> Dict[str, Any]:
     """ Process a single job. """
@@ -181,7 +187,7 @@ async def _process_job(config: Dict[str, Any], job: Dict[str, Any], hook) -> Dic
         if inspect.isgeneratorfunction(handler) or inspect.isasyncgenfunction(handler):
             log.debug("SLS Core | Running job as a generator.")
             generator_output = rp_job.run_job_generator(handler, job)
-            aggregated_output = {'output': []}
+            aggregated_output: dict[str, typing.Any] = {'output': []}
 
             async for part in generator_output:
                 log.debug(f"SLS Core | Streaming output: {part}", job['id'])
@@ -210,6 +216,7 @@ async def _process_job(config: Dict[str, Any], job: Dict[str, Any], hook) -> Dic
     finally:
         log.debug(f"SLS Core | Posting output: {result}", job['id'])
         hook.post_output(job['id'], result)
+        return result
 
 
 # ---------------------------------------------------------------------------- #