def ensure_stream_ingestion_jobs(client: feast_spark.Client,
all_projects: bool):
"""Ensures all required stream ingestion jobs are running and cleans up the unnecessary jobs.
More concretely, it will determine
- which stream ingestion jobs are running
- which stream ingestion jobs should be running
And it'll do 2 kinds of operations
- Cancel all running jobs that should not be running
- Start all non-existent jobs that should be running
Args:
all_projects (bool): If true, runs the check for all project.
Otherwise only checks the client's current project.
"""
projects = (client.feature_store.list_projects()
if all_projects else [client.feature_store.project])
expected_job_hash_to_table_refs = _get_expected_job_hash_to_table_refs(
client, projects)
expected_job_hashes = set(expected_job_hash_to_table_refs.keys())
jobs_by_hash: Dict[str, StreamIngestionJob] = {}
for job in client.list_jobs(include_terminated=False):
if isinstance(job, StreamIngestionJob):
jobs_by_hash[job.get_hash()] = job
existing_job_hashes = set(jobs_by_hash.keys())
job_hashes_to_cancel = existing_job_hashes - expected_job_hashes
job_hashes_to_start = expected_job_hashes - existing_job_hashes
logging.debug(
f"existing_job_hashes = {sorted(list(existing_job_hashes))} expected_job_hashes = {sorted(list(expected_job_hashes))}"
)
for job_hash in job_hashes_to_cancel:
job = jobs_by_hash[job_hash]
logging.info(
f"Cancelling a stream ingestion job with job_hash={job_hash} job_id={job.get_id()} status={job.get_status()}"
)
try:
job.cancel()
except FailedPrecondition as exc:
logging.warning(f"Job canceling failed with exception {exc}")
for job_hash in job_hashes_to_start:
# Any job that we wish to start should be among expected table refs map
project, table_name = expected_job_hash_to_table_refs[job_hash]
logging.info(
f"Starting a stream ingestion job for project={project}, table_name={table_name} with job_hash={job_hash}"
)
feature_table = client.feature_store.get_feature_table(name=table_name,
project=project)
client.start_stream_to_online_ingestion(feature_table, [],
project=project)
def test_list_jobs_long_table_name(
feast_client: Client,
feast_spark_client: SparkClient,
batch_source: Union[BigQuerySource, FileSource],
):
entity = Entity(name="long_entity_name" * 10,
description="S2id",
value_type=ValueType.INT64)
feature_table = FeatureTable(
name="just1a2featuretable3with4a5really6really7really8really9really10",
entities=[entity.name],
features=[Feature("unique_drivers", ValueType.INT64)],
batch_source=batch_source,
)
feast_client.apply(entity)
feast_client.apply(feature_table)
data_sample = generate_data().rename(columns={"s2id": entity.name})
feast_client.ingest(feature_table, data_sample)
job = feast_spark_client.start_offline_to_online_ingestion(
feature_table,
data_sample.event_timestamp.min().to_pydatetime(),
data_sample.event_timestamp.max().to_pydatetime() +
timedelta(seconds=1),
)
wait_retry_backoff(
lambda: (None, job.get_status() == SparkJobStatus.COMPLETED), 180)
all_job_ids = [
job.get_id() for job in feast_spark_client.list_jobs(
include_terminated=True,
project=feast_client.project,
table_name=feature_table.name,
)
]
assert job.get_id() in all_job_ids
features = feast_client.get_online_features(
[f"{feature_table.name}:unique_drivers"],
entity_rows=[{
entity.name: key
} for key in data_sample[entity.name].tolist()],
).to_dict()
ingested = pd.DataFrame.from_dict(features)
pd.testing.assert_frame_equal(
ingested[[entity.name, f"{feature_table.name}:unique_drivers"]],
data_sample[[entity.name, "unique_drivers"]].rename(
columns={"unique_drivers": f"{feature_table.name}:unique_drivers"
}),
)
def test_list_jobs_long_table_name(
feast_client: Client,
feast_spark_client: SparkClient,
batch_source: Union[BigQuerySource, FileSource],
):
entity = Entity(
name="s2id",
description="S2id",
value_type=ValueType.INT64,
)
feature_table = FeatureTable(
name=
"just1a2featuretable3with4a5really6really7really8really9really10really11really12long13name",
entities=["s2id"],
features=[Feature("unique_drivers", ValueType.INT64)],
batch_source=batch_source,
)
feast_client.apply(entity)
feast_client.apply(feature_table)
data_sample = generate_data()
feast_client.ingest(feature_table, data_sample)
job = feast_spark_client.start_offline_to_online_ingestion(
feature_table,
data_sample.event_timestamp.min().to_pydatetime(),
data_sample.event_timestamp.max().to_pydatetime() +
timedelta(seconds=1),
)
wait_retry_backoff(
lambda: (None, job.get_status() == SparkJobStatus.COMPLETED), 180)
all_job_ids = [
job.get_id() for job in feast_spark_client.list_jobs(
include_terminated=True,
project=feast_client.project,
table_name=feature_table.name,
)
]
assert job.get_id() in all_job_ids
def ensure_stream_ingestion_jobs(client: Client, all_projects: bool):
"""Ensures all required stream ingestion jobs are running and cleans up the unnecessary jobs.
More concretely, it will determine
- which stream ingestion jobs are running
- which stream ingestion jobs should be running
And it'll do 2 kinds of operations
- Cancel all running jobs that should not be running
- Start all non-existent jobs that should be running
Args:
all_projects (bool): If true, runs the check for all project.
Otherwise only checks the client's current project.
"""
projects = (client.feature_store.list_projects()
if all_projects else [client.feature_store.project])
if client.config.exists(opt.WHITELISTED_PROJECTS):
whitelisted_projects = client.config.get(opt.WHITELISTED_PROJECTS)
if whitelisted_projects:
whitelisted_projects = whitelisted_projects.split(",")
projects = [
project for project in projects
if project in whitelisted_projects
]
expected_job_hash_to_tables = _get_expected_job_hash_to_tables(
client, projects)
expected_job_hashes = set(expected_job_hash_to_tables.keys())
jobs_by_hash: Dict[str, StreamIngestionJob] = {}
# when we want to retry failed jobs, we shouldn't include terminated jobs here
# thus, Control Loop will behave like no job exists and will spawn new one
for job in client.list_jobs(include_terminated=not client.config.
getboolean(opt.JOB_SERVICE_RETRY_FAILED_JOBS)):
if (isinstance(job, StreamIngestionJob)
and job.get_status() != SparkJobStatus.COMPLETED):
jobs_by_hash[job.get_hash()] = job
existing_job_hashes = set(jobs_by_hash.keys())
job_hashes_to_cancel = existing_job_hashes - expected_job_hashes
job_hashes_to_start = expected_job_hashes - existing_job_hashes
logger.debug(f"existing_job_hashes = {sorted(list(existing_job_hashes))} "
f"expected_job_hashes = {sorted(list(expected_job_hashes))}")
for job_hash in job_hashes_to_start:
# Any job that we wish to start should be among expected table refs map
project, feature_table = expected_job_hash_to_tables[job_hash]
logger.warning(
f"Starting a stream ingestion job for project={project}, "
f"table_name={feature_table.name} with job_hash={job_hash}")
client.start_stream_to_online_ingestion(feature_table, [],
project=project)
# prevent scheduler from peak load
time.sleep(client.config.getint(opt.JOB_SERVICE_PAUSE_BETWEEN_JOBS))
for job_hash in job_hashes_to_cancel:
job = jobs_by_hash[job_hash]
if job.get_status() != SparkJobStatus.IN_PROGRESS:
logger.warning(
f"Can't cancel job with job_hash={job_hash} job_id={job.get_id()} status={job.get_status()}"
)
continue
logger.warning(
f"Cancelling a stream ingestion job with job_hash={job_hash} job_id={job.get_id()} status={job.get_status()}"
)
try:
job.cancel()
except FailedPrecondition as exc:
logger.error(f"Job canceling failed with exception {exc}")