Hey @brouberol, we are currently working on this task and there are some open questions.
Here is the idea:
We will build a kokkuri image which will include the logic for copying the trained model from the PVC to an S3 bucket.
This logic will be implemented in ml-piplines repo following the same logic and CI/CD processes with the retraining image. More precisely, this code will be containerised via kokkuri pipelines during the gitlab-ci process, and the generated image will be pushed to machine-learning/ml-pipelines/ docker-registry.
This image will be run via WMFKubernetesPodOperator in the tone-check retraining DAG as the last step of the pipeline.

Questions

• Do we need specific permissions for copying files from the PVC to an S3 bucket on the POD side and/or in the container itself?
• Do we need something specific when creating the S3 bucket that we will use for exporting the models?
• How do you find this architecture to build the model_export functionality on the side of ml-piplines repo and not in the Airflow-DAGs repo?

The branch that we are working is this one.

Do we need specific permissions for copying files from the PVC to an S3 bucket on the POD side and/or in the container itself?

No. The only thing you'd need would be a network policy allowing egress traffic to S3 as well as credentials stored in the airlfow-ml private values, injected in the Airflow connections, to authenticate access to S3.

Do we need something specific when creating the S3 bucket that we will use for exporting the models?

You'll need to have created a specific keypair to create the bucket with, and distribute this keypair to the Airflow private values, as previously stated. We (DPE SRE) can do that.

How do you find this architecture to build the model_export functionality on the side of ml-piplines repo and not in the Airflow-DAGs repo?

If it works for you, then all the better. If you think that this is something that could benefit everyone, then wmf_airflow_common might be a better place for it, but if it is ml-related and you're happy with it, then by all means.

Q: How do you find this architecture to build the model_export functionality on the side of ml-piplines repo and not in the Airflow-DAGs repo?
A: If it works for you, then all the better. If you think that this is something that could benefit everyone, then wmf_airflow_common might be a better place for it, but if it is ml-related and you're happy with it, then by all means.

Initially I was planning to implement it directly on the airflow-dags side, in order to use the get_s3_client() method from wmf_airflow_common.clients.s3 import get_s3_client.
The get_s3_client() is running via PythonOperator, can I attach the PVC on it and run the whole thing without the WMFKubernetesPodOperator ? This option is kinda easier by avoiding:

1. Build job logic in ml-pieplines
2. Build and publish the image via kokkuri
3. Set up the connections/keypairs etc...

So, if we can attach the PVC using the PythonOperator and use the get_s3_client(conn_id="thanos_swift") that would be easier I think.
I remember we had created this "thanos_swift" connection for a previous task.

Questions

• Can I attach the PVC using the PythonOperator in order to use the wmf_airflow_common.clients.s3 and do the job using the "thanos_swift" connection?
• Can I do the opposite? Use the wmf_airflow_common.clients.s3 directly in WMFKubernetesPodOperator ?

Can I attach the PVC using the PythonOperator in order to use the wmf_airflow_common.clients.s3 and do the job using the "thanos_swift" connection?

No, a PVC is a Kubernetes notion. A PythonOperator is just a way to execute a python function in your task pod. It has no way to attach and mount a Persistent Volume onto its runtime. If you need to read from Ceph, then your need a KubernetesPodOperator, with volume mounting from the ceph PVC.

Can I do the opposite? Use the wmf_airflow_common.clients.s3 directly in WMFKubernetesPodOperator ?

To do this, you'd need to have your KubernetesPodOperatorrun a pod that uses the airflow image itself, which should give you access to the wmf_airflow_common package. That being said, it would require the logic executed by your DAG to be fully defined in airflow-dags then, I think.

Can I do the opposite? Use the wmf_airflow_common.clients.s3 directly in WMFKubernetesPodOperator ?

To do this, you'd need to have your KubernetesPodOperatorrun a pod that uses the airflow image itself, which should give you access to the wmf_airflow_common package. That being said, it would require the logic executed by your DAG to be fully defined in airflow-dags then, I think.

This is perfect!
So is we use an image from: docker-registry/repos/data-engineering/airflow/ via the WMFKubernetesPodOperator, then we can just build the logic directly on the dags.
Thank you very much for your time @brouberol!

Yep, that should work!

Update

I moved all the code and logic into the Airflow-DAGs repo. I used three different airflow images and I had different issues:

1. Using airflow:ebd2d6ee4509445c05572ab5fa6e6afb588fb285-production or airflow:latest produced a ModuleNotFoundError: No module named 'wmf_airflow_common':

[2025-10-20, 13:35:51 UTC] {pod_manager.py:536} INFO - [base] === First list all elements in /mnt/model-training/tone_check/20251014T132011/output_model ===
[2025-10-20, 13:35:51 UTC] {pod_manager.py:536} INFO - [base] total 4
[2025-10-20, 13:35:51 UTC] {pod_manager.py:536} INFO - [base] drwxrwsr-x 2 app runuser 4096 Oct 15 16:59 checkpoint-63618
[2025-10-20, 13:35:51 UTC] {pod_manager.py:536} INFO - [base] === Checking model output path ===
[2025-10-20, 13:35:51 UTC] {pod_manager.py:536} INFO - [base] ✅ Model directory found. Uploading to S3...
[2025-10-20, 13:35:51 UTC] {pod_manager.py:536} INFO - [base] Traceback (most recent call last):
[2025-10-20, 13:35:51 UTC] {pod_manager.py:536} INFO - [base]   File "<stdin>", line 3, in <module>
[2025-10-20, 13:35:51 UTC] {pod_manager.py:555} INFO - [base] ModuleNotFoundError: No module named 'wmf_airflow_common'
[2025-10-20, 13:35:51 UTC] {pod.py:1122} INFO - Deleting pod: upload-model-to-s3-ebodqj6

2. Using airflow:2025-06-23-122527-798105179248020e79ffe44a8cb442d8675e9db7 tag it never loaded in the airflow UI (it was parsing the code for ever).

3. Using airflow-2.10.5-py3.11-2025-10-08-171529-e8280ad4b26af3cc41fd00af1ee6840463c2e0f4 it threw the exact same error.

From the above logs, we can see that the container starts correctly, prints the statements and runs until the wmf_airflow_common import.
So, it seems that it cannot find the wmf_common files. Do I miss something?
I tried to build the image locally run it interactively but I cannot see the wmf code. Do I need to pip install it or something?

I tried to avoid importing the s3_client using from wmf_airflow_common.clients.s3 import get_s3_client and implement the logic directly in the container like this.
Now I am getting this error:

[2025-10-22, 11:56:44 UTC] {pod_manager.py:536} INFO - [base] Unable to load the config, contains a configuration error.
[2025-10-22, 11:56:44 UTC] {pod_manager.py:536} INFO - [base] Traceback (most recent call last):
[2025-10-22, 11:56:44 UTC] {pod_manager.py:536} INFO - [base]   File "/usr/lib/python3.11/pathlib.py", line 1117, in mkdir
[2025-10-22, 11:56:44 UTC] {pod_manager.py:536} INFO - [base]     os.mkdir(self, mode)
[2025-10-22, 11:56:44 UTC] {pod_manager.py:536} INFO - [base] FileNotFoundError: [Errno 2] No such file or directory: '/opt/airflow/logs/scheduler/2025-10-22'
[2025-10-22, 11:56:44 UTC] {pod_manager.py:536} INFO - [base] During handling of the above exception, another exception occurred:

You can see the full logs in the following pase:

Update

After some discussions with people from the DE team, I am pasting here some ideas and good practices which answer the above comments.

Ideas from DE

The main problem seems to be (in the above comment) that we are missing the volumes and volumeMounts part of the airflow pod spec.
So we could add these to the task pod, but not sure that this is the most sustainable or most efficient way to achieve it.
Screenshots from the output of kubectl describe pod airflow-scheduler-7c6464cbdd-w8hzp :

We would probably be interested in the discussion happening at: T405360: Implement an Airflow operator for moving data from point A to B
In our case, we trying to move data from a locally mounted cephf's PVC, which just looks like a local filesystem, to an S3 bucket on swift/thanos.
We might find that using our data-engineering/sync-utils image image is efficient and we could use rclone or s3cmd to push to the file(s) to swift.

We shouldn't have to go about re-adding fundamental airflow infrastructure in the tasks, which is why I think that probably using an airflow image in a KubernetesPodOperator based task is probably over-complicated.

However, maybe there is another way:

Just as you can customise the pod that is launched by the KubernetesPodOperator.

• We can customise the pod that is launched by the KubernetesExecutor kubernetes_executor.html#pod-template-file

• This is better, we can use a pod_override for the KubernetesExecutor: kubernetes_executor.html#pod-override

This means that we can mount the PVC within the pod that runs the task.

We have some examples of where we modify the pod that is used to run the task:
Code Examples:

• data-engineering/airflow-dags/common.py
• data-engineering/airflow-dags/refine_to_hive_hourly_dag_factory.py

Both of the above examples are where we modify the resources associated with the executor pod that launches a KubernetesPodOperator based task. So there are two pods in this kind of setup.
In the current case, you could override the executor_config for the pod to mount the volume, then we could use the standard PythonOperator and get the S3 creds with wmf_airflow_common.clients.s3 as you had originally intended. No need to use a KubernetesPodOperator at all.

Update

I configured the executor_pod_override to mount the PVC so we can use it in our DAG with all the components (pure Airflow components and WMFKuberentesPodOperator).
I tested the final step move_model_to_s3_task using the get_s3_client() from wmf_airflow_common/clients/ and it was working perfectly.
The model was copied in the dedicated S3 bucket for exporting retrained models "s3://wmf-ml-models/retrained-models".

$ s3cmd -c /etc/s3cmd/cfg.d/ml-team.cfg ls -H --recursive s3://wmf-ml-models/retrained-models/tone-check/checkpoint-63618/
2026-01-20 13:33   865   s3://wmf-ml-models/retrained-models/tone-check/checkpoint-63618/config.json
2026-01-20 13:33   678M  s3://wmf-ml-models/retrained-models/tone-check/checkpoint-63618/model.safetensors
2026-01-20 13:33  1357M  s3://wmf-ml-models/retrained-models/tone-check/checkpoint-63618/optimizer.pt
2026-01-20 13:33    13K  s3://wmf-ml-models/retrained-models/tone-check/checkpoint-63618/rng_state.pth
2026-01-20 13:33  1064   s3://wmf-ml-models/retrained-models/tone-check/checkpoint-63618/scheduler.pt
2026-01-20 13:33   695   s3://wmf-ml-models/retrained-models/tone-check/checkpoint-63618/special_tokens_map.json
2026-01-20 13:33     2M  s3://wmf-ml-models/retrained-models/tone-check/checkpoint-63618/tokenizer.json
2026-01-20 13:33  1330   s3://wmf-ml-models/retrained-models/tone-check/checkpoint-63618/tokenizer_config.json
2026-01-20 13:33    24K  s3://wmf-ml-models/retrained-models/tone-check/checkpoint-63618/trainer_state.json
2026-01-20 13:33     5K  s3://wmf-ml-models/retrained-models/tone-check/checkpoint-63618/training_args.bin
2026-01-20 13:33   972K  s3://wmf-ml-models/retrained-models/tone-check/checkpoint-63618/vocab.txt

You can find the logs of the move_model_to_s3_task in the following paste: