11import logging
2+ import os
3+ import signal
4+ import threading
5+ import time
6+ import traceback
27from concurrent .futures import ThreadPoolExecutor
8+ from typing import Dict , List , Tuple
39
410import grpc
511
612import feast
13+ from feast .constants import CONFIG_JOB_SERVICE_ENABLE_CONTROL_LOOP
714from feast .core import JobService_pb2_grpc
815from feast .core .JobService_pb2 import (
916 CancelJobResponse ,
3138)
3239from feast .pyspark .launcher import (
3340 get_job_by_id ,
41+ get_stream_to_online_ingestion_params ,
3442 list_jobs ,
3543 start_historical_feature_retrieval_job ,
3644 start_offline_to_online_ingestion ,
4351)
4452
4553
54+ def _job_to_proto (spark_job : SparkJob ) -> JobProto :
55+ job = JobProto ()
56+ job .id = spark_job .get_id ()
57+ status = spark_job .get_status ()
58+ if status == SparkJobStatus .COMPLETED :
59+ job .status = JobStatus .JOB_STATUS_DONE
60+ elif status == SparkJobStatus .IN_PROGRESS :
61+ job .status = JobStatus .JOB_STATUS_RUNNING
62+ elif status == SparkJobStatus .FAILED :
63+ job .status = JobStatus .JOB_STATUS_ERROR
64+ elif status == SparkJobStatus .STARTING :
65+ job .status = JobStatus .JOB_STATUS_PENDING
66+ else :
67+ raise ValueError (f"Invalid job status { status } )
68+
69+ if isinstance (spark_job , RetrievalJob ):
70+ job .type = JobType .RETRIEVAL_JOB
71+ job .retrieval .output_location = spark_job .get_output_file_uri (block = False )
72+ elif isinstance (spark_job , BatchIngestionJob ):
73+ job .type = JobType .BATCH_INGESTION_JOB
74+ elif isinstance (spark_job , StreamIngestionJob ):
75+ job .type = JobType .STREAM_INGESTION_JOB
76+ else :
77+ raise ValueError (f"Invalid job type { job } )
78+
79+ return job
80+
81+
4682class JobServiceServicer (JobService_pb2_grpc .JobServiceServicer ):
47- def __init__ (self ):
48- self .client = feast .Client ()
49-
50- def _job_to_proto (self , spark_job : SparkJob ) -> JobProto :
51- job = JobProto ()
52- job .id = spark_job .get_id ()
53- status = spark_job .get_status ()
54- if status == SparkJobStatus .COMPLETED :
55- job .status = JobStatus .JOB_STATUS_DONE
56- elif status == SparkJobStatus .IN_PROGRESS :
57- job .status = JobStatus .JOB_STATUS_RUNNING
58- elif status == SparkJobStatus .FAILED :
59- job .status = JobStatus .JOB_STATUS_ERROR
60- elif status == SparkJobStatus .STARTING :
61- job .status = JobStatus .JOB_STATUS_PENDING
62- else :
63- raise ValueError (f"Invalid job status { status } )
64-
65- if isinstance (spark_job , RetrievalJob ):
66- job .type = JobType .RETRIEVAL_JOB
67- job .retrieval .output_location = spark_job .get_output_file_uri (block = False )
68- elif isinstance (spark_job , BatchIngestionJob ):
69- job .type = JobType .BATCH_INGESTION_JOB
70- elif isinstance (spark_job , StreamIngestionJob ):
71- job .type = JobType .STREAM_INGESTION_JOB
72- else :
73- raise ValueError (f"Invalid job type { job } )
74-
75- return job
83+ def __init__ (self , client ):
84+ self .client = client
7685
7786 def StartOfflineToOnlineIngestionJob (
7887 self , request : StartOfflineToOnlineIngestionJobRequest , context
@@ -117,6 +126,20 @@ def StartStreamToOnlineIngestionJob(
117126 feature_table = self .client .get_feature_table (
118127 request .table_name , request .project
119128 )
129+
130+ if self .client ._config .getboolean (CONFIG_JOB_SERVICE_ENABLE_CONTROL_LOOP ):
131+ # If the control loop is enabled, return existing stream ingestion job id instead of starting a new one
132+ params = get_stream_to_online_ingestion_params (
133+ self .client , request .project , feature_table , []
134+ )
135+ job_hash = params .get_job_hash ()
136+ for job in list_jobs (include_terminated = True , client = self .client ):
137+ if isinstance (job , StreamIngestionJob ) and job .get_hash () == job_hash :
138+ return StartStreamToOnlineIngestionJobResponse (id = job .get_id ())
139+ raise RuntimeError (
140+ "Feast Job Service has control loop enabled, but couldn't find the existing stream ingestion job for the given FeatureTable"
141+ )
142+
120143 # TODO: add extra_jars to request
121144 job = start_stream_to_online_ingestion (
122145 client = self .client ,
@@ -131,7 +154,7 @@ def ListJobs(self, request, context):
131154 jobs = list_jobs (
132155 include_terminated = request .include_terminated , client = self .client
133156 )
134- return ListJobsResponse (jobs = [self . _job_to_proto (job ) for job in jobs ])
157+ return ListJobsResponse (jobs = [_job_to_proto (job ) for job in jobs ])
135158
136159 def CancelJob (self , request , context ):
137160 """Stop a single job"""
@@ -142,7 +165,30 @@ def CancelJob(self, request, context):
142165 def GetJob (self , request , context ):
143166 """Get details of a single job"""
144167 job = get_job_by_id (request .job_id , client = self .client )
145- return GetJobResponse (job = self ._job_to_proto (job ))
168+ return GetJobResponse (job = _job_to_proto (job ))
169+
170+
171+ def start_control_loop () -> None :
172+ """Starts control loop that continuously ensures that correct jobs are being run.
173+
174+ Currently this affects only the stream ingestion jobs. Please refer to
175+ ensure_stream_ingestion_jobs for full documentation on how the check works.
176+
177+ """
178+ logging .info (
179+ "Feast Job Service is starting a control loop in a background thread, "
180+ "which will ensure that stream ingestion jobs are successfully running."
181+ )
182+ try :
183+ client = feast .Client ()
184+ while True :
185+ ensure_stream_ingestion_jobs (client , all_projects = True )
186+ time .sleep (1 )
187+ except Exception :
188+ traceback .print_exc ()
189+ finally :
190+ # Send interrupt signal to the main thread to kill the server if control loop fails
191+ os .kill (os .getpid (), signal .SIGINT )
146192
147193
148194class HealthServicer (HealthService_pb2_grpc .HealthServicer ):
@@ -156,18 +202,110 @@ def intercept_service(self, continuation, handler_call_details):
156202 return continuation (handler_call_details )
157203
158204
159- def start_job_service ():
205+ def start_job_service () -> None :
160206 """
161207 Start Feast Job Service
162208 """
163209
164210 log_fmt = "%(asctime)s %(levelname)s %(message)s"
165211 logging .basicConfig (level = logging .INFO , format = log_fmt )
166212
213+ client = feast .Client ()
214+
215+ if client ._config .getboolean (CONFIG_JOB_SERVICE_ENABLE_CONTROL_LOOP ):
216+ # Start the control loop thread only if it's enabled from configs
217+ thread = threading .Thread (target = start_control_loop , daemon = True )
218+ thread .start ()
219+
167220 server = grpc .server (ThreadPoolExecutor (), interceptors = (LoggingInterceptor (),))
168- JobService_pb2_grpc .add_JobServiceServicer_to_server (JobServiceServicer (), server )
221+ JobService_pb2_grpc .add_JobServiceServicer_to_server (
222+ JobServiceServicer (client ), server
223+ )
169224 HealthService_pb2_grpc .add_HealthServicer_to_server (HealthServicer (), server )
170225 server .add_insecure_port ("[::]:6568" )
171226 server .start ()
172- print ("Feast job server listening on port :6568" )
227+ logging . info ("Feast Job Service is listening on port :6568" )
173228 server .wait_for_termination ()
229+
230+
231+ def _get_expected_job_hash_to_table_refs (
232+ client : feast .Client , projects : List [str ]
233+ ) -> Dict [str , Tuple [str , str ]]:
234+ """
235+ Checks all feature tables for the requires project(s) and determines all required stream
236+ ingestion jobs from them. Outputs a map of the expected job_hash to a tuple of (project, table_name).
237+
238+ Args:
239+ all_projects (bool): If true, runs the check for all project.
240+ Otherwise only checks the current project.
241+
242+ Returns:
243+ Dict[str, Tuple[str, str]]: Map of job_hash -> (project, table_name) for expected stream ingestion jobs
244+ """
245+ job_hash_to_table_refs = {}
246+
247+ for project in projects :
248+ feature_tables = client .list_feature_tables (project )
249+ for feature_table in feature_tables :
250+ if feature_table .stream_source is not None :
251+ params = get_stream_to_online_ingestion_params (
252+ client , project , feature_table , []
253+ )
254+ job_hash = params .get_job_hash ()
255+ job_hash_to_table_refs [job_hash ] = (project , feature_table .name )
256+
257+ return job_hash_to_table_refs
258+
259+
260+ def ensure_stream_ingestion_jobs (client : feast .Client , all_projects : bool ):
261+ """Ensures all required stream ingestion jobs are running and cleans up the unnecessary jobs.
262+
263+ More concretely, it will determine
264+ - which stream ingestion jobs are running
265+ - which stream ingestion jobs should be running
266+ And it'll do 2 kinds of operations
267+ - Cancel all running jobs that should not be running
268+ - Start all non-existent jobs that should be running
269+
270+ Args:
271+ all_projects (bool): If true, runs the check for all project.
272+ Otherwise only checks the client's current project.
273+ """
274+
275+ projects = client .list_projects () if all_projects else [client .project ]
276+
277+ expected_job_hash_to_table_refs = _get_expected_job_hash_to_table_refs (
278+ client , projects
279+ )
280+
281+ expected_job_hashes = set (expected_job_hash_to_table_refs .keys ())
282+
283+ jobs_by_hash : Dict [str , StreamIngestionJob ] = {}
284+ for job in client .list_jobs (include_terminated = False ):
285+ if isinstance (job , StreamIngestionJob ):
286+ jobs_by_hash [job .get_hash ()] = job
287+
288+ existing_job_hashes = set (jobs_by_hash .keys ())
289+
290+ job_hashes_to_cancel = existing_job_hashes - expected_job_hashes
291+ job_hashes_to_start = expected_job_hashes - existing_job_hashes
292+
293+ logging .debug (
294+ f"existing_job_hashes = { sorted (list (existing_job_hashes ))} { sorted (list (expected_job_hashes ))}
295+ )
296+
297+ for job_hash in job_hashes_to_cancel :
298+ job = jobs_by_hash [job_hash ]
299+ logging .info (
300+ f"Cancelling a stream ingestion job with job_hash={ job_hash } { job .get_id ()} { job .get_status ()}
301+ )
302+ job .cancel ()
303+
304+ for job_hash in job_hashes_to_start :
305+ # Any job that we wish to start should be among expected table refs map
306+ project , table_name = expected_job_hash_to_table_refs [job_hash ]
307+ logging .info (
308+ f"Starting a stream ingestion job for project={ project } { table_name } { job_hash }
309+ )
310+ feature_table = client .get_feature_table (name = table_name , project = project )
311+ client .start_stream_to_online_ingestion (feature_table , [], project = project )
0 commit comments