Use Colab as a base image. #1444

djherbis · 2024-11-27T03:29:55Z

This change makes a number of major changes:

Colab is the base image
uv is the main package install tool
leveraging requirements.txt instead of many separate installs
stop building and installing tensorflow/torch/lightbgm/jax since those are managed by the Colab base image now

In order to decide what packages to explicitly install I:

looked at what packages are in the Colab base image
looked at what packages were in the Kaggle image
looked at what packages were explicitly mentioned in Kaggle Dockerfile

This may still take a few iterations to get all the right parts in the image, but this should hopefully make the image much more manageable.

http://b/365782129

This change makes a number of major changes: - Colab is the base image - uv is the main package install tool - leveraging requirements.txt instead of many separate installs - stop building and installing tensorflow/torch/lightbgm/jax since those are managed by the Colab base image now In order to decide what packages to explicitly install I: - looked at what packages are in the Colab base image - looked at what packages were in the Kaggle image - looked at what packages were explicitly mentioned in Kaggle Dockerfile This may still take a few iterations to get all the right parts in the image, but this should hopefully make the image much more manageable. http://b/365782129

djherbis · 2024-11-27T21:35:36Z

Jenkinsfile

-          }
-        }
-      }
-    }


These are no longer needed since they are installed in the base image.

djherbis · 2024-11-27T21:36:41Z

test

@@ -3,7 +3,7 @@ set -e

 IMAGE_TAG='kaggle/python-build'
 IMAGE_TAG_OVERRIDE=''
-ADDITONAL_OPTS=''
+ADDITONAL_OPTS='--runtime runc ' # Use the CPU runtime by default


Colab image fails if using the nvidia runtime (which Jenkins does by default) while on a CPU VM.

djherbis · 2024-11-27T21:37:25Z

tests/test_fastai.py

        learn = tabular_learner(dls, layers=[200, 100])
-        learn.fit_one_cycle(n_epoch=1)
+        with learn.no_bar():


fastprogress was trying to render html while running in the background and it caused issues, so I turned off the progress bar rendering.

rosbo

Awesome work on this!

rosbo · 2024-11-27T21:57:00Z

config.txt

 CUDA_MAJOR_VERSION=12
-CUDA_MINOR_VERSION=3
+CUDA_MINOR_VERSION=2


I assume this downgrade is to align with Colab?

Yes! Colab is on 12.2, so we are too now. Sucks to go backwards, but better that we're aligned.

rosbo · 2024-11-27T21:58:44Z

tests/test_lightgbm.py

@@ -34,7 +34,9 @@ def test_cpu(self):

        self.assertEqual(1, gbm.best_iteration)

+    # TODO(b/381256047): Colab needs to install GPU-enabled lightgbm.


Do you consider this a launch blocker for this?

Given how popular lightgbm (on GPU) is, I'm planning to have it be a launch blocker.
I'm working with the Colab team to get it in there :)

re-added it in #1451

Colab is also working on adding it to the base so we can re-drop once its out in their image

calderjo · 2024-12-02T18:24:25Z

Dockerfile.tmpl

+# b/183041606#comment5: the Kaggle data proxy doesn't support these APIs. If the library is missing, it falls back to using a regular BigQuery query to fetch data.
+RUN uv pip uninstall --system google-cloud-bigquery-storage
+
+# NOTE(herbison): uv fails to install this for some reason


prefer a bug over ldap, unless you intend to change it to a uv install soon.

I'll fix this in #1448.

calderjo · 2024-12-02T18:24:36Z

Dockerfile.tmpl

+# NOTE(herbison): uv fails to install this for some reason
+RUN pip install git+https://github.com/Kaggle/learntools
+
+# We install an incompatible pair of libs (shapely<, libpysal==4.9.2) so we can't put this one in the requirements.txt


you mention that it can't be added, but it's still in the requirements.txt:

https://github.com/Kaggle/docker-python/pull/1444/files#diff-2fb70e5756d0842b40d8bbe8b41a8be4d3a68c7d98fdaac2914b1dbf43638ed3R75-R76

did you intend on removing it?

also be nice to drop the related bug number here too: b/328788268

Its not fully pinned in the requirements.txt, its a "<=4.9.2" there, but then we "upgrade it to "==4.9.2" here.

calderjo · 2024-12-02T18:25:13Z

Dockerfile.tmpl

-    /tmp/clean-layer.sh
-{{ end}}
-
-# Install PyTorch


calderjo · 2024-12-02T18:26:18Z

kaggle_requirements.txt

@@ -0,0 +1,139 @@
+altair>=5.4.0


super optional: seem sorta alphabetically ordered, is this intentional?
do we want to add a comment at the top to enforce it? ngl it is easier to read.

Yeah, I originally created this list by doing pip freeze on Colab and on Kaggle, then diffing the results and taking the diff subset, and then only including explicitly mentioned packages from our Dockerfile.

It is roughly alpha ordered because of pip freeze (caps is going above other things, caps probably shouldn't matter, I'm not sure why some packages are capitalized). I'll add a comment in my follow-up.

calderjo · 2024-12-02T18:26:39Z

kaggle_requirements.txt

+vtk
+wandb
+wavio
+xgboost==2.0.3


maybe include " b/350573866: xgboost v2.1.0 breaks learntools " comment from before

Agreed, I've been adding back pins and bug info in #1448

francesco-ballarin · 2024-12-07T09:43:26Z

Hi @djherbis ,
a couple of questions:

where do I find the built docker image after this PR? If I pull gcr.io/kaggle-images/python I still get an image where python3 comes from conda, which I guess is the image that was built just before this PR.
now that you are based on the colab image, can you consider publishing backend info (basically a dump of apt list and a dump of pip freeze) just like colab does in https://github.com/googlecolab/backend-info ?

Thanks,
Francesco

djherbis added 12 commits November 27, 2024 03:25

rm micromamba since its not actually applying to the env

c4ee63d

fix kagglehub

292110d

keep Colab key deps

dc18661

fix gcloud config path

ffabed3

use cpu runtime for cpu tests

e458e59

rm lightgbm comment

44d2385

remove unused docker run cmd

a62abe1

match ipywidgets version thats used now

c5fb602

skip lightgbm test until underlying bug fixed

0104ac6

disable cuML test since its not on colab

d916094

add ctx bug for disable cuML

89f13e6

djherbis requested review from calderjo and rosbo November 27, 2024 21:34

djherbis commented Nov 27, 2024

View reviewed changes

rosbo approved these changes Nov 27, 2024

View reviewed changes

djherbis merged commit 66bac48 into main Nov 27, 2024
4 checks passed

djherbis deleted the uv-colab-base-image branch November 27, 2024 22:05

calderjo reviewed Dec 2, 2024

View reviewed changes

Dockerfile.tmpl

/tmp/clean-layer.sh

{{ end}}

# Install PyTorch

Copy link

Contributor

calderjo Dec 2, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🥳🎉

calderjo reviewed Dec 2, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Use Colab as a base image. #1444

Use Colab as a base image. #1444

djherbis commented Nov 27, 2024

djherbis Nov 27, 2024

djherbis Nov 27, 2024

djherbis Nov 27, 2024

rosbo left a comment

rosbo Nov 27, 2024

djherbis Nov 27, 2024

rosbo Nov 27, 2024

djherbis Nov 27, 2024

djherbis Dec 4, 2024

calderjo Dec 2, 2024

djherbis Dec 2, 2024 •

edited

Loading

calderjo Dec 2, 2024 •

edited

Loading

djherbis Dec 2, 2024

calderjo Dec 2, 2024

calderjo Dec 2, 2024

djherbis Dec 2, 2024

calderjo Dec 2, 2024 •

edited

Loading

djherbis Dec 2, 2024

francesco-ballarin commented Dec 7, 2024

		@@ -34,7 +34,9 @@ def test_cpu(self):

		self.assertEqual(1, gbm.best_iteration)

		# TODO(b/381256047): Colab needs to install GPU-enabled lightgbm.

Use Colab as a base image. #1444

Use Colab as a base image. #1444

Conversation

djherbis commented Nov 27, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rosbo left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

djherbis Dec 2, 2024 • edited Loading

Choose a reason for hiding this comment

calderjo Dec 2, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

calderjo Dec 2, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

francesco-ballarin commented Dec 7, 2024

djherbis Dec 2, 2024 •

edited

Loading

calderjo Dec 2, 2024 •

edited

Loading

calderjo Dec 2, 2024 •

edited

Loading