Skip to content

Commit fbb5eb8

Browse files
pyalexpyalex
committed
solve techdebt: move communication with job to JobCoordinator
1 parent e343c2f commit fbb5eb8

File tree

11 files changed

+457
-299
lines changed

11 files changed

+457
-299
lines changed

‎core/src/main/java/feast/core/config/FeastProperties.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,10 @@ public static class FeatureSetSpecStreamProperties {
172172

173173
/* Kafka topic to receive acknowledgment from ingestion job on successful processing of new specs */
174174
@NotBlank private String specsAckTopic = "feast-feature-set-specs-ack";
175+
176+
/* Notify jobs interval in millisecond.
177+
How frequently Feast will check on Pending FeatureSets and publish them to kafka. */
178+
@Positive private long notifyIntervalMilliseconds;
175179
}
176180
}
177181

‎core/src/main/java/feast/core/dao/FeatureSetRepository.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package feast.core.dao;
1818

1919
import feast.core.model.FeatureSet;
20+
import feast.proto.core.FeatureSetProto;
2021
import java.util.List;
2122
import org.springframework.data.jpa.repository.JpaRepository;
2223

@@ -37,4 +38,7 @@ public interface FeatureSetRepository extends JpaRepository<FeatureSet, String>
3738
// find all feature sets matching the given name pattern and project pattern
3839
List<FeatureSet> findAllByNameLikeAndProject_NameLikeOrderByNameAsc(
3940
String name, String project_name);
41+
42+
// find all feature sets matching given status
43+
List<FeatureSet> findAllByStatus(FeatureSetProto.FeatureSetStatus status);
4044
}

‎core/src/main/java/feast/core/job/JobUpdateTask.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import feast.proto.core.FeatureSetProto;
2525
import java.time.Instant;
2626
import java.util.List;
27+
import java.util.Map;
2728
import java.util.Optional;
2829
import java.util.concurrent.Callable;
2930
import java.util.concurrent.ExecutionException;
@@ -32,6 +33,7 @@
3233
import java.util.concurrent.Future;
3334
import java.util.concurrent.TimeUnit;
3435
import java.util.concurrent.TimeoutException;
36+
import java.util.stream.Collectors;
3537
import lombok.Getter;
3638
import lombok.extern.slf4j.Slf4j;
3739
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.hash.Hashing;
@@ -149,10 +151,24 @@ private Job startJob(String jobId) {
149151
}
150152

151153
private void updateFeatureSets(Job job) {
154+
Map<FeatureSet, FeatureSetJobStatus> alreadyConnected =
155+
job.getFeatureSetJobStatuses().stream()
156+
.collect(Collectors.toMap(FeatureSetJobStatus::getFeatureSet, s -> s));
157+
152158
for (FeatureSet fs : featureSets) {
159+
if (alreadyConnected.containsKey(fs)) {
160+
continue;
161+
}
162+
153163
FeatureSetJobStatus status = new FeatureSetJobStatus();
154164
status.setFeatureSet(fs);
155165
status.setJob(job);
166+
if (fs.getStatus() == FeatureSetProto.FeatureSetStatus.STATUS_READY) {
167+
// Feature Set was already delivered to previous generation of the job
168+
// (another words, it exists in kafka)
169+
// so we expect Job will ack latest version based on history from kafka topic
170+
status.setVersion(fs.getVersion());
171+
}
156172
status.setDeliveryStatus(FeatureSetProto.FeatureSetJobDeliveryStatus.STATUS_IN_PROGRESS);
157173
job.getFeatureSetJobStatuses().add(status);
158174
}
@@ -175,6 +191,7 @@ private Job updateStatus(Job job) {
175191
}
176192

177193
job.setStatus(newStatus);
194+
updateFeatureSets(job);
178195
return job;
179196
}
180197

‎core/src/main/java/feast/core/model/FeatureSetJobStatus.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
*/
1717
package feast.core.model;
1818

19+
import com.google.common.base.Objects;
1920
import feast.proto.core.FeatureSetProto.FeatureSetJobDeliveryStatus;
2021
import java.io.Serializable;
2122
import javax.persistence.*;
@@ -61,4 +62,23 @@ public FeatureSetJobStatusKey() {}
6162
@Enumerated(EnumType.STRING)
6263
@Column(name = "delivery_status")
6364
private FeatureSetJobDeliveryStatus deliveryStatus;
65+
66+
@Column(name = "version", columnDefinition = "integer default 0")
67+
private int version;
68+
69+
@Override
70+
public boolean equals(Object o) {
71+
if (this == o) return true;
72+
if (o == null || getClass() != o.getClass()) return false;
73+
FeatureSetJobStatus that = (FeatureSetJobStatus) o;
74+
return version == that.version
75+
&& Objects.equal(job.getId(), that.job.getId())
76+
&& Objects.equal(featureSet.getReference(), that.featureSet.getReference())
77+
&& deliveryStatus == that.deliveryStatus;
78+
}
79+
80+
@Override
81+
public int hashCode() {
82+
return Objects.hashCode(job, featureSet, deliveryStatus, version);
83+
}
6484
}

‎core/src/main/java/feast/core/service/JobCoordinatorService.java‎

Lines changed: 135 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616
*/
1717
package feast.core.service;
1818

19+
import static feast.core.model.FeatureSet.parseReference;
20+
1921
import com.google.protobuf.InvalidProtocolBufferException;
2022
import feast.core.config.FeastProperties;
2123
import feast.core.config.FeastProperties.JobProperties;
@@ -26,21 +28,24 @@
2628
import feast.core.model.*;
2729
import feast.proto.core.CoreServiceProto.ListStoresRequest.Filter;
2830
import feast.proto.core.CoreServiceProto.ListStoresResponse;
31+
import feast.proto.core.FeatureSetProto;
32+
import feast.proto.core.IngestionJobProto;
2933
import feast.proto.core.StoreProto;
3034
import feast.proto.core.StoreProto.Store.Subscription;
3135
import java.util.ArrayList;
3236
import java.util.HashSet;
3337
import java.util.List;
3438
import java.util.Optional;
3539
import java.util.Set;
36-
import java.util.concurrent.ExecutionException;
37-
import java.util.concurrent.ExecutorCompletionService;
38-
import java.util.concurrent.ExecutorService;
39-
import java.util.concurrent.Executors;
40+
import java.util.concurrent.*;
4041
import java.util.stream.Collectors;
4142
import javax.validation.constraints.Positive;
4243
import lombok.extern.slf4j.Slf4j;
44+
import org.apache.kafka.clients.consumer.ConsumerRecord;
4345
import org.springframework.beans.factory.annotation.Autowired;
46+
import org.springframework.data.util.Pair;
47+
import org.springframework.kafka.annotation.KafkaListener;
48+
import org.springframework.kafka.core.KafkaTemplate;
4449
import org.springframework.scheduling.annotation.Scheduled;
4550
import org.springframework.stereotype.Service;
4651
import org.springframework.transaction.annotation.Transactional;
@@ -49,24 +54,29 @@
4954
@Service
5055
public class JobCoordinatorService {
5156

57+
private final int SPEC_PUBLISHING_TIMEOUT_SECONDS = 5;
58+
5259
private final JobRepository jobRepository;
5360
private final FeatureSetRepository featureSetRepository;
5461
private final SpecService specService;
5562
private final JobManager jobManager;
5663
private final JobProperties jobProperties;
64+
private final KafkaTemplate<String, FeatureSetProto.FeatureSetSpec> specPublisher;
5765

5866
@Autowired
5967
public JobCoordinatorService(
6068
JobRepository jobRepository,
6169
FeatureSetRepository featureSetRepository,
6270
SpecService specService,
6371
JobManager jobManager,
64-
FeastProperties feastProperties) {
72+
FeastProperties feastProperties,
73+
KafkaTemplate<String, FeatureSetProto.FeatureSetSpec> specPublisher) {
6574
this.jobRepository = jobRepository;
6675
this.featureSetRepository = featureSetRepository;
6776
this.specService = specService;
6877
this.jobManager = jobManager;
6978
this.jobProperties = feastProperties.getJobs();
79+
this.specPublisher = specPublisher;
7080
}
7181

7282
/**
@@ -153,4 +163,124 @@ public Optional<Job> getJob(Source source, Store store) {
153163
// return the latest
154164
return Optional.of(jobs.get(0));
155165
}
166+
167+
@Transactional
168+
@Scheduled(fixedDelayString = "${feast.stream.specsOptions.notifyIntervalMilliseconds}")
169+
public void notifyJobsWhenFeatureSetUpdated() {
170+
List<FeatureSet> pendingFeatureSets =
171+
featureSetRepository.findAllByStatus(FeatureSetProto.FeatureSetStatus.STATUS_PENDING);
172+
173+
pendingFeatureSets.stream()
174+
.filter(
175+
fs -> {
176+
List<FeatureSetJobStatus> runningJobs =
177+
fs.getJobStatuses().stream()
178+
.filter(jobStatus -> jobStatus.getJob().isRunning())
179+
.collect(Collectors.toList());
180+
181+
return runningJobs.size() > 0
182+
&& runningJobs.stream()
183+
.allMatch(jobStatus -> jobStatus.getVersion() < fs.getVersion());
184+
})
185+
.forEach(
186+
fs -> {
187+
log.info("Sending new FeatureSet {} to Ingestion", fs.getReference());
188+
189+
// Sending latest version of FeatureSet to all currently running IngestionJobs
190+
// (there's one topic for all sets).
191+
// All related jobs would apply new FeatureSet on the fly.
192+
// In case kafka doesn't respond within SPEC_PUBLISHING_TIMEOUT_SECONDS we will try
193+
// again later.
194+
try {
195+
specPublisher
196+
.sendDefault(fs.getReference(), fs.toProto().getSpec())
197+
.get(SPEC_PUBLISHING_TIMEOUT_SECONDS, TimeUnit.SECONDS);
198+
} catch (Exception e) {
199+
log.error(
200+
"Error occurred while sending FeatureSetSpec to kafka. Cause {}."
201+
+ " Will retry later",
202+
e.getMessage());
203+
return;
204+
}
205+
206+
// Updating delivery status for related jobs (that are currently using this
207+
// FeatureSet).
208+
// We now set status to IN_PROGRESS, so listenAckFromJobs would be able to
209+
// monitor delivery progress for each new version.
210+
fs.getJobStatuses().stream()
211+
.filter(s -> s.getJob().isRunning())
212+
.forEach(
213+
jobStatus -> {
214+
jobStatus.setDeliveryStatus(
215+
FeatureSetProto.FeatureSetJobDeliveryStatus.STATUS_IN_PROGRESS);
216+
jobStatus.setVersion(fs.getVersion());
217+
});
218+
featureSetRepository.saveAndFlush(fs);
219+
});
220+
}
221+
222+
/**
223+
* Listener for ACK messages coming from IngestionJob when FeatureSetSpec is installed (in
224+
* pipeline).
225+
*
226+
* <p>Updates FeatureSetJobStatus for respected FeatureSet (selected by reference) and Job (select
227+
* by Id).
228+
*
229+
* <p>When all related (running) to FeatureSet jobs are updated - FeatureSet receives READY status
230+
*
231+
* @param record ConsumerRecord with key: FeatureSet reference and value: Ack message
232+
*/
233+
@KafkaListener(topics = {"${feast.stream.specsOptions.specsAckTopic}"})
234+
@Transactional
235+
public void listenAckFromJobs(
236+
ConsumerRecord<String, IngestionJobProto.FeatureSetSpecAck> record) {
237+
String setReference = record.key();
238+
Pair<String, String> projectAndSetName = parseReference(setReference);
239+
FeatureSet featureSet =
240+
featureSetRepository.findFeatureSetByNameAndProject_Name(
241+
projectAndSetName.getSecond(), projectAndSetName.getFirst());
242+
if (featureSet == null) {
243+
log.warn(
244+
String.format("ACKListener received message for unknown FeatureSet %s", setReference));
245+
return;
246+
}
247+
248+
int ackVersion = record.value().getFeatureSetVersion();
249+
250+
if (featureSet.getVersion() != ackVersion) {
251+
log.warn(
252+
String.format(
253+
"ACKListener received outdated ack for %s. Current %d, Received %d",
254+
setReference, featureSet.getVersion(), ackVersion));
255+
return;
256+
}
257+
258+
log.info("Updating featureSet {} delivery statuses.", featureSet.getReference());
259+
260+
featureSet.getJobStatuses().stream()
261+
.filter(
262+
js ->
263+
js.getJob().getId().equals(record.value().getJobName())
264+
&& js.getVersion() == ackVersion)
265+
.findFirst()
266+
.ifPresent(
267+
featureSetJobStatus ->
268+
featureSetJobStatus.setDeliveryStatus(
269+
FeatureSetProto.FeatureSetJobDeliveryStatus.STATUS_DELIVERED));
270+
271+
boolean allDelivered =
272+
featureSet.getJobStatuses().stream()
273+
.filter(js -> js.getJob().isRunning())
274+
.allMatch(
275+
js ->
276+
js.getDeliveryStatus()
277+
.equals(FeatureSetProto.FeatureSetJobDeliveryStatus.STATUS_DELIVERED));
278+
279+
if (allDelivered) {
280+
log.info("FeatureSet {} update is completely delivered", featureSet.getReference());
281+
282+
featureSet.setStatus(FeatureSetProto.FeatureSetStatus.STATUS_READY);
283+
featureSetRepository.saveAndFlush(featureSet);
284+
}
285+
}
156286
}

0 commit comments

Comments
 (0)