damccorm commented on code in PR #39074:
URL: https://github.com/apache/beam/pull/39074#discussion_r4048131931
##########
sdks/python/apache_beam/ml/inference/base.py:
##########
@@ -406,6 +408,225 @@ def should_garbage_collect_on_timeout(self) -> bool:
return self.share_model_across_processes()
+class SubprocessModelHandler(ModelHandler[ExampleT, PredictionT, ModelT], ABC):
+ """Base class for model handlers that spin up a subprocess server."""
+ @abstractmethod
+ def get_port(self, model: ModelT) -> int:
+ """Returns the port the subprocess server is listening on."""
+ pass
+
+ @abstractmethod
+ def get_model_name(self) -> str:
+ """Returns the model name."""
+ pass
+
+ @abstractmethod
+ def check_connectivity(self, model: ModelT) -> None:
+ """Checks connectivity to the server and attempts to recover/mark for
restart."""
+ pass
+
+
+class SubProcessModelServer:
+ """Manages the lifecycle of a generic subprocess model server."""
+ def __init__(self, handler_path: str, model_name: str, port: int = None,
temp_dir: tempfile.TemporaryDirectory = None):
+ self._handler_path = handler_path
+ self._model_name = model_name
+ self._port = port
+ self._temp_dir = temp_dir
+ self._process = None
+ self._server_started = False
+ self._server_process_lock = threading.RLock()
+ self.start_server()
+
+ def start_server(self, retries=3):
+ with self._server_process_lock:
+ if not self._server_started:
+ if self._process:
+ logging.info("Terminating existing generic subprocess model server
before restart")
+ try:
+ self._process.terminate()
+ self._process.wait(timeout=5)
+ except Exception:
+ try:
+ self._process.kill()
+ except Exception:
+ pass
+ self._process = None
+ self._port = None
+
+ from apache_beam.utils import subprocess_server
+ if self._port is None:
+ self._port, = subprocess_server.pick_port(None)
+
+ cmd = [
+ sys.executable,
+ '-m',
+ 'apache_beam.ml.inference.subprocess_server',
+ '--handler_path',
+ self._handler_path,
+ '--port',
+ str(self._port),
+ ]
+ logging.info("Starting generic model server with %s", cmd)
+ self._process = subprocess.Popen(
+ cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+
+ # Emit the output of this command as info level logging.
+ def log_stdout():
+ line = self._process.stdout.readline()
+ while line:
+ logging.info(line.decode(errors='backslashreplace').rstrip())
+ line = self._process.stdout.readline()
+
+ t = threading.Thread(target=log_stdout)
+ t.daemon = True
+ t.start()
+
+ self.check_connectivity(retries)
+
+ def get_server_port(self) -> int:
+ if not self._server_started:
+ self.start_server()
+ return self._port
+
+ def check_connectivity(self, retries=3):
+ import urllib.request
+ import urllib.error
+
+ url = f"http://localhost:{self._port}/v1/models";
+ attempts = 0
+ max_attempts = 12 # 12 * 5s = 60s timeout
+ while self._process.poll() is None and attempts < max_attempts:
+ try:
+ # Use standard library to check connectivity to avoid extra
dependencies
+ req = urllib.request.Request(url, method="GET")
+ with urllib.request.urlopen(req, timeout=5) as response:
+ if response.status == 200:
+ self._server_started = True
+ return
+ except urllib.error.URLError:
+ pass
+ except Exception as e:
+ logging.warning("Error checking connectivity: %s", e)
+ attempts += 1
+ time.sleep(5)
+
+ if retries == 0:
+ self._server_started = False
+ raise Exception(
+ "Failed to start generic subprocess server, polling process exited
with code " +
+ f"{self._process.poll()}. Next time a request is tried, the server
will be restarted"
+ )
+ else:
+ self.start_server(retries - 1)
Review Comment:
Fixed. Wrapped check_connectivity with self._server_process_lock and reset
self._server_started = False before calling start_server.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
