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}")