Support nested image spec (#3284)

Signed-off-by: Akinori Mitani <akinori@artera.ai>

Author: amitani

flytekit/image_spec/default_builder.py

@@ -86,9 +86,12 @@
 id=micromamba \
     --mount=from=micromamba,source=/usr/bin/micromamba,target=/usr/bin/micromamba \
     micromamba config set use_lockfiles False && \
-    micromamba create -n runtime --root-prefix /opt/micromamba \
+    ( micromamba create -n runtime --root-prefix /opt/micromamba \
     -c conda-forge $CONDA_CHANNELS \
-    python=$PYTHON_VERSION $CONDA_PACKAGES
+    python=$PYTHON_VERSION $CONDA_PACKAGES \
+    || micromamba install -n runtime --root-prefix /opt/micromamba \
+    -c conda-forge $CONDA_CHANNELS \
+    python=$PYTHON_VERSION $CONDA_PACKAGES )
 """)
 
 DOCKER_FILE_TEMPLATE = Template("""\
@@ -98,14 +101,15 @@
 
 FROM $BASE_IMAGE
 
+WORKDIR /
 USER root
 $APT_INSTALL_COMMAND
 RUN --mount=from=micromamba,source=/etc/ssl/certs/ca-certificates.crt,target=/tmp/ca-certificates.crt \
     [ -f /etc/ssl/certs/ca-certificates.crt ] || \
     mkdir -p /etc/ssl/certs/ && cp /tmp/ca-certificates.crt /etc/ssl/certs/ca-certificates.crt
 
 RUN id -u flytekit || useradd --create-home --shell /bin/bash flytekit
-RUN chown -R flytekit /root && chown -R flytekit /home
+RUN chown flytekit /root && chown -R flytekit /home
 
 $INSTALL_PYTHON_TEMPLATE
 

tests/flytekit/unit/core/image_spec/test_image_spec.py

@@ -93,6 +93,29 @@ def test_image_spec(mock_image_spec_builder, monkeypatch):
     assert image_spec.commands == ["echo hello"]
 
 
+@pytest.mark.skipif(
+    os.environ.get("_FLYTEKIT_TEST_IMAGE_BUILD_ENGINE", "0") == "0",
+    reason="Set _FLYTEKIT_TEST_IMAGE_BUILD_ENGINE=1 to run this test",
+)
+def test_nested_build(monkeypatch):
+    monkeypatch.setenv("FLYTE_PUSH_IMAGE_SPEC", "0")
+    base_image = ImageSpec(
+        name="base",
+        packages=["torch"],
+        python_version="3.11"
+    )
+
+    image_spec = ImageSpec(
+        name="final",
+        packages=["pandas"],
+        python_version="3.11",
+        base_image=base_image,
+    )
+    assert image_spec._is_force_push is False
+
+    ImageBuildEngine.build(image_spec)
+
+
 def test_image_spec_engine_priority():
     new_image_name = "fqn.xyz/flytekit"
     mock_image_builder_10 = Mock()