-
Notifications
You must be signed in to change notification settings - Fork 185
Proposal - Supporting user-provided ML models in workflows #6616
What we're proposing
Let users bring their own trained ML models into Texera and use them in workflows, upload a model, keep versions of it, share it, and run inference on your data — just like datasets work today. Tracking issue: #6494.
What learned from research
Studied how other data/ML platforms handle this, and took a few clear lessons:
- Each resource type gets its own table and its own resolver. Don't force one shared table or one universal path rule to cover everything — give models their own home, separate from datasets.
- Tell resource types apart by a label in the path, then look up the right table from that label (e.g.
/datasets/...vs/models/...). - A model is just a folder of files with versions — the same shape as a dataset. So we can reuse the storage and versioning.
- To run a model, give the worker a real folder of its files (via a mount), because model loaders expect to open files from a directory.
The design
Models are a sibling of datasets: their own tables and their own storage repo, on top of the same storage engine we already use. The path label decides the type (/datasets/... vs /models/...), and a worker mounts the chosen model version to read its files.
Tables and storage
texera-asset-tables-design
How a model reaches a worker when it runs
texera-geesefs-model-mount
The plan
One small step per PR (see sub-tasks under #6494): add the path label → model tables → model storage & lookup → upload/version/access API → UI → use in a workflow → sharing & docs. We start with PyTorch; more frameworks and a no-code inference operator come later.
Questions for the community
- Path label: we're making the label required (a leading
datasets/modelssegment), with a one-time migration of existing dataset paths. Any concerns with requiring it, versus leaving datasets with no label as the default? - Separate
modeltables (mirroring the dataset tables) vs. reusing the dataset tables with a type column — any preference or downside we're missing? - Loading models by mounting the files into the worker — anyone have experience or concerns?
Feedback welcome!
All reactions
Replies: 6 comments 7 replies
@tanishqgandhi1908 Thanks for starting the discussion. Our protocol is to create issues for a discussion, and raise PRs for each issue. Can you organize the content in this way?
All reactions
Design Update
Key Decisions
- Isolate mounts per computing unit (CU)
- Context: We currently support read-only operations for model inference, which allows using shared volumes and mounting points across CUs.
- Decision: We are opting for dedicated per-CU dataset mounts.
- Rationale: This architecture prepares the system for upcoming write support (model training), ensuring two users working with the same base model can modify their mounts independently without conflict.
- Offload mounting from CU pods to a privileged per-node DaemonSet
- Context: GeeseFS requires elevated privileges to perform mounts. Granting these directly to CU pods introduces security risks, as users can execute arbitrary code that could compromise adjacent pods on the node.
- Decision: A dedicated per-node DaemonSet handles the mounts and attaches the resulting volumes to the CU pods.
- Defer dataset resolution and mounting to runtime
- Decision: Mount execution entirely to runtime.
- Rationale: Mounts are only required while a workflow is actively executing so we do not want to perform mounting before or after running.
- Integrate model selection into the operator property panel
- Decision: Use a specialized UI parameter kind built on @carloea2's implementation (feat(udf): enable Python UDF UI parameters end to end #5912 ) rather than introducing a parallel selection mechanism:
- Details: The parameter resolves directly to the local directory where the specific model version is mounted.
- Provide contextual sample code based on model and storage types
- Decision: Auto-generate tailored Python UDF snippets based on the selected model type and storage backing to accelerate user implementation.
Under Research
Next area of focus—design still in progress:
- Hardware Constraints (GPU & VRAM): CUs currently enforce GPU allocation limits. We are defining model-side metadata requirements and validation checks to prevent GPU-dependent models from running on unallocated CUs.
- OS-Level Portability: Investigating how tightly coupled trained models are to their original host OS environment and defining cross-platform execution boundaries.
All reactions
End to end user experience
Continuing on the above discussion, please find below the end-to-end user experience. Models can be created, versioned, shared and used in a workflow. Here's what the flow looks like for a user.
- Create a model
Same as creating a dataset — name it, and pick a framework (PyTorch, TensorFlow, ONNX, scikit-learn, or "other") and a format (TorchScript, safetensors, ONNX, SavedModel, joblib, pickle, ...). Both can be changed later.
-
Upload the files
Drag a folder in and the structure is kept. Big weight files upload in chunks, and models get their own size limit — 2 GB per file by default, where datasets are 20 MB. Every upload becomes a new version you can go back to. -
Share it
Invite people by email with read or write access, or make it public. There's a separate switch for whether people can download it. Models also show up in search — by name, description, framework or format — and public ones appear in the Hub with view and like counts. -
Copy the code to use it
Each model page has a Usage tab that writes the Python for you, based on the framework and format you picked. torch.jit.load for TorchScript, joblib.load for scikit-learn, an ONNX Runtime session for ONNX, and so on — with comments explaining each step.
- Paste it and pick a version
The code declares its own parameter, so pasting it into a Python UDF makes a row appear in the property panel with a model picker. Choose a version there and it reaches your code as a local folder — already mounted, nothing to download.
Would love suggestions on this flow- anything you'd expect to be here and isn't.
All reactions
@mengw15 @bobbai00 : Please check this comment #6616 (comment) about this "DaemonSet". I think it's related to a recent change you made about having a separate micro-service to manage the Postgres credential to make sure each pod is safe. If so, please chime in and give your thoughts.
All reactions
At high level, it is indeed related to #4242 . Specifically, the motivation and the solution are shared across two problems.
However, I did see a major difference here:
- In Postgres-credential case, CU communicates using Restful API over HTTP protocol
- In ML model mounting case, CU uses POSIX file I/O over FUSE protocol, not a network protocol.
Therefore, design wise:
- Postgres-credential case uses a micro service
- ML-model case uses something like DaemonSet.
I am not quite familiar with DaemonSet. But I want to propose few questions that the final design needs to answer:
- What damage can a malicious user cause over other users' models ? In postgres-credential case, a malicious user cannot cause any damages to other users' Iceberg tables, because the access is controlled using JWT token.
- Is the solution general? Let's say, a CU is running on an EC2 instance. Postgres-credential solution is general because the a CU on EC2 can still talk to a micro service running somewhere else. However, for DaemonSet case, the the CU has to be running as a K8s pod and co-located with a daemon set pod. This seems to indicate that the solution is not general.
I hope @tanishqgandhi1908 and @aicam can answer these two questions. I hope there is a design that can limit the damage and be general.
All reactions
Agree with @bobbai00's two questions — the key property to preserve from the Postgres-credential design is that credentials never enter the CU pod and access is scoped per user; the DaemonSet mounts should provide the same guarantee.
All reactions
Thanks @bobbai00 , regarding your questions
1- To ensure security, we create a S3 proxy in file-service (check file-service/src/main/scala/org/apache/texera/service/util/S3ProxyServlet.scala), and GeeseFS thinks our S3 proxy is the exact LakeFS but in fact, we authorize requests, GeeseFS sends user JWT token as S3_access_secret since it only support S3 authorization, and we use it to determine if user has access to the model
2- Yes, we ran a test and deployed application on local Kubernetes and used S3 bucket on AWS and everything worked. But it was much slower which is expected due to network bandwidth.
All reactions
I asked the input from @bobbai00 and @mengw15 to verify the high-level similarity between this "DaemonSet" and the micro-service we introduced to keep Postgres credentials. I am glad you two verified my claim.
@aicam : Is "DaemonSet" also micro service? If not, why shouldn't it be a micro service?
@bobbai00 : Please verify if you agree with the answers of @aicam to your two questions.
All reactions
texera-mounter is a microservice, DaemonSet is a kind in K8s
All reactions
@aicam Your comment is not clear. Can you elaborate on the differences?
All reactions
In Kubernetes (K8s), calling both a "microservice" is fine from an architectural standpoint, but their operational behavior and scheduling strategies in a cluster are fundamentally different.
-
Node Placement & Scheduling Rules
DaemonSet: Kubernetes automatically schedules exactly one pod on every single node (or a targeted subset of nodes matching specific labels/selectors). If you have 5 nodes, you get 5 instances—one per node.
Standard Microservice (Deployment): The K8s Scheduler decides where pods land based on available CPU/memory, node affinity, or anti-affinity rules. One node might run three instances, while another node runs zero. -
Scaling Dynamics
DaemonSet: Scales implicitly with your infrastructure. When a new physical or virtual node is added to the cluster, K8s immediately provisions a DaemonSet pod onto that new node without manual intervention.
Standard Microservice (Deployment): Scales explicitly based on application traffic or workload demand (e.g., via a Horizontal Pod Autoscaler based on CPU usage or custom metrics), completely independent of how many nodes exist in the cluster. -
Network Routing & Locality
DaemonSet: Often uses hostNetwork: true or hostPort to attach directly to the local node's host interface. Applications on that node usually reach out to localhost to talk to the local daemon instance.
Standard Microservice (Deployment): Sits behind a Kubernetes Service with an internal ClusterIP. Incoming requests are round-robined across any available pod in the cluster regardless of which node it is running on.
For example in Texera, we have config-service as a Deployment, it sits behind Envoy Gateway, receives requests from users and connects to Postgres but does not have any access to node network or operating system. In contrast, mounter which is a DaemonSet, it does not receive request from public users (only from access-control service), does not connect to Postgres but it has full access to node file system and can modify K8s volumes.
All reactions
-
👍 1
@parshimers @mengw15 @Ma77Ball reviewed the design, and security was raised as a primary concern. To address this, a flag was added to disable 'mounting' by default. Even if a user enables the model tab in admin settings, the mounting feature stays disabled until it's enabled via the environment variable.