Add options to disable concurrent image pulling and saving squashfs files to the dedicated file/directory, config to expose enroot logs

[itechdima]

Description

Added new srun argument:

• --container-image-save - (PATH) to indicate, where you can save squashfs file of your image (could be either file or directory). Squashfs files are not deleted after srun, following srun calls with the same argument value will reuse existing squashfs file

For this option it's possible to configure default behaviour through pyxis configuration. Also, extended configuration with option:

• expose_enroot_logs - (BOOLEAN) changes enroot logs from the in-memory file to stderr, so now every pyxis step is visible.

Tests

Performed 2 runs for different setups on same slurm cluster.

1. pyxis 0.20:
   • first run: https://gist.github.com/itechdima/9f97cf6e71d847ec7c4a19f3c8ce1e43

     105 tests, 13 failures, 1 skipped
     StartTime=2025-01-27T11:45:49 EndTime=2025-01-27T12:31:21
     

   • second run: https://gist.github.com/itechdima/a1319ce649e8efa13575548d07015ca8

     105 tests, 13 failures, 1 skipped
     StartTime=2025-01-27T12:38:31 EndTime=2025-01-27T13:26:43
     

2. patched pyxis without any additional options:
   • first run: https://gist.github.com/itechdima/70cf20bbb61134f7801f3cde42271376

     105 tests, 11 failures, 1 skipped
     StartTime=2025-01-28T09:46:29 EndTime=2025-01-28T10:33:50
     

   • second run: https://gist.github.com/itechdima/385c0e77b4934711c181e796d588ab9a

     105 tests, 11 failures, 1 skipped
     StartTime=2025-01-28T11:22:17 EndTime=2025-01-28T12:07:19
     

   Note: squashfs_image.bats tests where fixed later, see 3d option and 3d run
3. patched pyxis with --container-image-shared and --container-image-save options in config:
   • first run: https://gist.github.com/itechdima/f9dc24a87c8cdcf6d45f37b3d9b9a6a9

     105 tests, 12 failures, 1 skipped
     StartTime=2025-01-27T14:34:40 EndTime=2025-01-27T15:10:13
     

   • second run: https://gist.github.com/itechdima/bba91cf76d212ff694ecb5a4ebafd418

     105 tests, 12 failures, 1 skipped
     StartTime=2025-01-27T15:11:40 EndTime=2025-01-27T15:33:42
     

   • third run (after fixes): https://gist.github.com/itechdima/c06a4a679ca2381d3f4dcfa7703775ed

     105 tests, 9 failures, 1 skipped
     StartTime=2025-01-30T14:19:49 EndTime=2025-01-30T14:43:08
     

I'm not sure about all the tests failures. Maybe some of them are flacky, and some of them are related to the:
"Some tests assume a specific enroot configuration (such as PMIx/PyTorch hooks), so they might not pass on all systems"
but overall, changes speed up even tests (instead of ~40min, all the test runs, except of the first one, will take ~20min)

Documentation

I didn't find a way to update github wiki with new pyxis config parameters, maybe you can help me with that.

[flx42]

I like the idea of being able to cache a squashfs in a shared location, but I think this is a bit too complex for a starting point, I will need to think a little more about the problem but here is my early feedback:

• expose_enroot_logs is not really related to the main problem you are trying to solve, so it should be removed from this PR.
• I think the feature should only be exposed on the command-line for now, not in the pyxis config file.
• Given the previous point, the cache folder should stay entirely under the user's control and we should not have pyxis create per-user folders automatically. Each user would specify a directory of their choice.

[itechdima]

@flx42
Thanks for the early reply (I wan't even able to complete my PR description before you started reviewing it :D )

expose_enroot_logs is not really related to the main problem you are trying to solve, so it should be removed from this PR.

sure, I can move it to a separate PR

I think the feature should only be exposed on the command-line for now, not in the pyxis config file.

The reason we decided to add not only to the command-line, but also to the pyxis config file, is that we can easily update config file for the whole Slurm cluster, and every existing script will automatically work with new functionality without any changes. If we remove it from the config, then it'll require every Slurm user to adapt their scripts (as for our case, we already have plenty of them and I think they would prefer it implicitly turned on in pyxis config, rather then try to understand where they need to specify that options)

Given the previous point, the cache folder should stay entirely under the user's control and we should not have pyxis create per-user folders automatically. Each user would specify a directory of their choice.

I see your point. Maybe it's better to change like it's done with pyxis runtime? Create only subdir with uid?

[itechdima]

@flx42
I've made some changes
Now, if the container image save path passed in the config, then the logic is similar to the pyxis runtime path: it doesn't create the whole path but only uid folder for the particular user.
If the container image save path passed through the srun argument, then it doesn't create any folder in that path

If you're OK with that, I'd like to discuss how would you prefer this PR to be split, so it'll be easier for you to review in more details :)

[itechdima]

@flx42 do you have any updates :)

[flx42]

Apologies, with GTC it's been a bit crazy here, will try to take a look at it soon.

[asteny]

@flx42 Hello Felix!
Can you please check PR of my colleague @itechdima ? We (in Nebius) want to use this feature in Soperator from NVIDIA upstream, but right now we are using forked code.

[flx42]

@flx42 Hello Felix! Can you please check PR of my colleague @itechdima ? We (in Nebius) want to use this feature in Soperator from NVIDIA upstream, but right now we are using forked code.

After discussing with internal users that have similar but also slightly different requirements, I will likely implement a different approach in pyxis in the future.

[flx42]

Here is the feature I have added: https://github.com/NVIDIA/pyxis/tree/main/importers
Let me know if that would work for your use case, or if you see something missing with this API, it can still evolve.

[flx42]

Closing as I'm not planning to merge this MR, please do open new issues if the importers mechanism is insufficient for your use case or if you find any bug.