feat: KEP-2655: Add data cache system by akshaychitneni · Pull Request #2755 · kubeflow/trainer

akshaychitneni · 2025-07-28T15:13:01Z

What this PR does / why we need it:
This PR adds distributed data cache as described in https://docs.google.com/document/d/1xj3K6bOT4f0EPiC4zwr2OsbRkROjIA8u_TBRZRVxrHI
Followup PRs will include SDK and trainer integrations

Proposal - kubeflow/community#864
KEP - #2655

Which issue(s) this PR fixes (optional, in Fixes #<issue number>, #<issue number>, ... format, will close the issue(s) when PR gets merged):
Fixes # #2757

Checklist:

Docs included if any changes are user facing

coveralls · 2025-07-28T15:26:37Z

Pull Request Test Coverage Report for Build 17837270829

Details

0 of 0 changed or added relevant lines in 0 files are covered.
No unchanged relevant lines lost coverage.
Overall coverage remained the same at 55.137%

Totals
Change from base Build 17772927799:	0.0%
Covered Lines:	1084
Relevant Lines:	1966

💛 - Coveralls

andreyvelich · 2025-08-11T14:57:39Z

+arrow = "55.0.0"
+arrow-schema = "55.0.0"


Shall we migrate to Arrow 55.2 as @comphead suggested ?

I am using datafusion 47.0.0 and it builds on arrow 55.0.0 - https://github.com/apache/datafusion/blob/e4433049b04ca2c1e2031eb05d1a0990210f11d6/Cargo.toml#L90. Let me try to address as a followup which might require additional validation

@akshaychitneni Can you create tracking issue for that please ?
We can address it in the followup PRs

andreyvelich · 2025-08-11T15:29:19Z

+#[derive(Serialize, Deserialize, Debug)]
+struct IndexPair {
+    start: u64,
+    end: u64,
+}


Can we use this struct from main.rs ?

andreyvelich · 2025-08-11T15:31:22Z

+
+    pub async fn init(&mut self) -> Result<()> {
+        let _ = self.fetch_data_files().await;
+        let df = self.ctx.sql("select * from memtable").await?.collect().await?;


Does DataFusion create the memtable for us ?

we register a table with name memtable here - https://github.com/kubeflow/trainer/pull/2755/files#diff-d6c85a17f636003dada6db08feb3052bf87b12f6672723af5b503ec390d9c980R94. Probably I will use the var instead of using name directly.

andreyvelich

@akshaychitneni I left a few comments.
Overall looks good.

Can you fix the CI please ?

/cc @kubeflow/kubeflow-trainer-team @rudeigerc @raravena80 @comphead
in case you want to provide any additional comments.

andreyvelich · 2025-09-14T18:18:10Z

+arrow = "55.0.0"
+arrow-schema = "55.0.0"


@akshaychitneni Can you create tracking issue for that please ?
We can address it in the followup PRs

google-oss-prow · 2025-09-14T18:29:12Z

@andreyvelich: GitHub didn't allow me to request PR reviews from the following users: kubeflow/kubeflow-trainer-team, raravena80.

Note that only kubeflow members and repo collaborators can review this PR, and authors cannot review their own PRs.

Details

In response to this:

@akshaychitneni I left a few comments.
Overall looks good.

Can you fix the CI please ?

/cc @kubeflow/kubeflow-trainer-team @rudeigerc @raravena80 @comphead
in case you want to provide any additional comments.

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository.

andreyvelich · 2025-09-14T18:30:10Z

@akshaychitneni Please can you also add a new scope to the PR title verification, e.g. cache, so we can do feat(cache):...: https://github.com/kubeflow/trainer/blob/master/.github/workflows/check-pr-title.yaml#L28

+[[package]]
+name = "tracing-subscriber"
+version = "0.3.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e8189decb5ac0fa7bc8b96b7cb9b2701d60d48805aca84a238004d665fcc4008"
+dependencies = [
+ "nu-ansi-term",
+ "sharded-slab",
+ "smallvec",
+ "thread_local",
+ "tracing-core",
+ "tracing-log",
+]


andreyvelich · 2025-09-18T16:30:14Z

@@ -0,0 +1,51 @@
+use std::env;


@akshaychitneni Could you also move main files for head and worker to the cmd Similar to controller manger: https://github.com/kubeflow/trainer/tree/master/cmd/trainer-controller-manager ?
I would suggest this structure:

cmd/cache/Dockerfile cmd/cache/head/main.rs cmd/cache/worker/main.rs

WDYT ?

However, I am not sure if Cargo supports such folder structure:

trainer/pkg/data_cache/Cargo.toml

Lines 27 to 33 in 492063d

[[bin]]

name = "head"

path = "cmd/head/main.rs"

[[bin]]

name = "worker"

path = "cmd/worker/main.rs"

This might not be helpful as it involves moving outside of cargo

andreyvelich

Thanks for this awesome effort 🚀
/lgtm
/assign @astefanutti @tenzen-y @Electronic-Waste @kubeflow/wg-training-leads @kubeflow/kubeflow-trainer-team
Please let us know if you have any additional comments before move this forward.

andreyvelich · 2025-09-18T17:47:46Z

/lgtm

Signed-off-by: Akshay Chitneni <achitneni@apple.com>

andreyvelich · 2025-09-19T00:41:03Z

/lgtm

andreyvelich

We can move forward with this PR and address the remaining comments in the followup changes.
Thanks again for this huge effort @akshaychitneni!
/approve

google-oss-prow · 2025-09-22T20:06:31Z

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: andreyvelich

The full list of commands accepted by this bot can be found here.

The pull request process is described here

Details

Needs approval from an approver in each of these files:

~~OWNERS~~ [andreyvelich]

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

Signed-off-by: Akshay Chitneni <achitneni@apple.com> Co-authored-by: Akshay Chitneni <achitneni@apple.com>

google-oss-prow Bot added the size/XXL label Jul 28, 2025

google-oss-prow Bot requested review from astefanutti and kuizhiqing July 28, 2025 15:13

akshaychitneni force-pushed the cache-oss branch from 9f5e926 to f82b6e2 Compare July 28, 2025 15:13

akshaychitneni changed the title ~~KEP-2655: Add data cache system~~ feat - KEP-2655: Add data cache system Jul 28, 2025

comphead reviewed Aug 5, 2025

View reviewed changes

andreyvelich reviewed Aug 6, 2025

View reviewed changes

Comment thread pkg/data_cache/README.md

Comment thread pkg/data_cache/README.md Outdated

akshaychitneni force-pushed the cache-oss branch 2 times, most recently from 3b7359b to b640ffa Compare August 8, 2025 03:27

andreyvelich reviewed Aug 11, 2025

View reviewed changes

akshaychitneni force-pushed the cache-oss branch 3 times, most recently from bff2ac3 to a3e89ae Compare August 12, 2025 23:41

akshaychitneni force-pushed the cache-oss branch from a3e89ae to e472c01 Compare August 24, 2025 16:45

andreyvelich reviewed Sep 14, 2025

View reviewed changes

google-oss-prow Bot requested review from comphead and rudeigerc September 14, 2025 18:29

andreyvelich added this to the v2.1 milestone Sep 17, 2025

akshaychitneni changed the title ~~feat - KEP-2655: Add data cache system~~ feat(cache) - KEP-2655: Add data cache system Sep 17, 2025

akshaychitneni force-pushed the cache-oss branch from e472c01 to 082f592 Compare September 17, 2025 23:39

akshaychitneni temporarily deployed to gpu-tests September 17, 2025 23:39 — with GitHub Actions Inactive

github-advanced-security AI found potential problems Sep 17, 2025

View reviewed changes

akshaychitneni changed the title ~~feat(cache) - KEP-2655: Add data cache system~~ feat(cache): KEP-2655: Add data cache system Sep 17, 2025

akshaychitneni mentioned this pull request Sep 18, 2025

KEP-2655: Data cache: Use Arrow 55.2 #2833

Closed

andreyvelich mentioned this pull request Sep 18, 2025

[Scheduling & Orchestration] Enhanced Batch Job Management - Part 1 kubernetes-sigs/ai-conformance#13

Open

akshaychitneni force-pushed the cache-oss branch from 082f592 to 492063d Compare September 18, 2025 16:19

akshaychitneni temporarily deployed to gpu-tests September 18, 2025 16:20 — with GitHub Actions Inactive

andreyvelich reviewed Sep 18, 2025

View reviewed changes

google-oss-prow Bot assigned andreyvelich Sep 18, 2025

google-oss-prow Bot added the lgtm label Sep 18, 2025

google-oss-prow Bot temporarily deployed to gpu-tests September 18, 2025 17:27 Inactive

akshaychitneni force-pushed the cache-oss branch from 492063d to 4f20335 Compare September 18, 2025 17:41

google-oss-prow Bot removed the lgtm label Sep 18, 2025

akshaychitneni temporarily deployed to gpu-tests September 18, 2025 17:41 — with GitHub Actions Inactive

andreyvelich changed the title ~~feat(cache): KEP-2655: Add data cache system~~ feat: KEP-2655: Add data cache system Sep 18, 2025

google-oss-prow Bot added the lgtm label Sep 18, 2025

Adding Data Cache

607de09

Signed-off-by: Akshay Chitneni <achitneni@apple.com>

akshaychitneni force-pushed the cache-oss branch from 4f20335 to 607de09 Compare September 18, 2025 18:00

google-oss-prow Bot removed the lgtm label Sep 18, 2025

akshaychitneni temporarily deployed to gpu-tests September 18, 2025 18:00 — with GitHub Actions Inactive

google-oss-prow Bot added the lgtm label Sep 19, 2025

google-oss-prow Bot temporarily deployed to gpu-tests September 19, 2025 00:46 Inactive

andreyvelich reviewed Sep 22, 2025

View reviewed changes

google-oss-prow Bot added the approved label Sep 22, 2025

google-oss-prow Bot merged commit bbaca24 into kubeflow:master Sep 22, 2025
34 of 35 checks passed

andreyvelich mentioned this pull request Oct 20, 2025

KEP-2655: Publish the Data Cache Implementation #2757

Closed

alexxfan pushed a commit to red-hat-data-services/trainer that referenced this pull request Nov 24, 2025

feat: KEP-2655: Add data cache system (kubeflow#2755)

938c1c3

Signed-off-by: Akshay Chitneni <achitneni@apple.com> Co-authored-by: Akshay Chitneni <achitneni@apple.com>

	[[bin]]
	name = "head"
	path = "cmd/head/main.rs"

	[[bin]]
	name = "worker"
	path = "cmd/worker/main.rs"

Conversation

akshaychitneni commented Jul 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

coveralls commented Jul 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pull Request Test Coverage Report for Build 17837270829

Details

💛 - Coveralls

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

andreyvelich left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

google-oss-prow Bot commented Sep 14, 2025

Uh oh!

andreyvelich commented Sep 14, 2025

Uh oh!

Check notice

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

andreyvelich left a comment

Choose a reason for hiding this comment

Uh oh!

andreyvelich commented Sep 18, 2025

Uh oh!

andreyvelich commented Sep 19, 2025

Uh oh!

andreyvelich left a comment

Choose a reason for hiding this comment

Uh oh!

google-oss-prow Bot commented Sep 22, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

akshaychitneni commented Jul 28, 2025 •

edited

Loading

coveralls commented Jul 28, 2025 •

edited

Loading