changes

Author: savingoyal

devtools/Makefile

@@ -250,8 +250,7 @@ shell: setup-tilt
 				"$$user_shell" -i; \
 		fi'
 
-# TODO: This can be done away with in a while since we now have metaflow-dev command.
-#
+
 # @echo '$(MAKE_CMD) create-dev-shell' >> $(DEVTOOLS_DIR)/start.sh
 # @echo 'rm -f /tmp/metaflow-devshell-*' >> $(DEVTOOLS_DIR)/start.sh
 create-dev-shell: setup-tilt

devtools/Tiltfile

@@ -99,6 +99,7 @@ if "minio" in enabled_components:
         set=[
             'rootUser=rootuser',
             'rootPassword=rootpass123',
+            # TODO: perturb the bucket name to avoid conflicts
             'buckets[0].name=metaflow-test',
             'buckets[0].policy=none',
             'buckets[0].purge=false',
@@ -600,6 +601,18 @@ if "ui" in enabled_components:
     config_resources.append('metaflow-ui')
     config_resources.append('metaflow-ui-static')
 
+cmd = '''
+ARCH=$(kubectl get nodes -o jsonpath='{.items[0].status.nodeInfo.architecture}')
+case "$ARCH" in
+  arm64)   echo linux-aarch64 ;;
+  amd64)   echo linux-64 ;;
+  *)       echo linux-64 ;;
+esac
+'''
+
+# For @conda/@pypi emulation
+metaflow_config["METAFLOW_KUBERNETES_CONDA_ARCH"] = str(local(cmd)).strip()
+
 local_resource(
     name="generate-configs",
     cmd=write_config_files(),

metaflow/metaflow_config.py

@@ -381,6 +381,8 @@
 # Default kubernetes QoS class
 KUBERNETES_QOS = from_conf("KUBERNETES_QOS", "burstable")
 
+# Architecture of kubernetes nodes - used for @conda/@pypi in metaflow-dev
+KUBERNETES_CONDA_ARCH = from_conf("KUBERNETES_CONDA_ARCH")
 ARGO_WORKFLOWS_KUBERNETES_SECRETS = from_conf("ARGO_WORKFLOWS_KUBERNETES_SECRETS", "")
 ARGO_WORKFLOWS_ENV_VARS_TO_SKIP = from_conf("ARGO_WORKFLOWS_ENV_VARS_TO_SKIP", "")
 

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -29,6 +29,7 @@
     KUBERNETES_SHARED_MEMORY,
     KUBERNETES_TOLERATIONS,
     KUBERNETES_QOS,
+    KUBERNETES_CONDA_ARCH,
 )
 from metaflow.plugins.resources_decorator import ResourcesDecorator
 from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
@@ -158,7 +159,7 @@ class KubernetesDecorator(StepDecorator):
 
     # Conda environment support
     supports_conda_environment = True
-    target_platform = "linux-aarch64"
+    target_platform = KUBERNETES_CONDA_ARCH or "linux-64"
 
     def init(self):
         super(KubernetesDecorator, self).init()

metaflow/plugins/pypi/conda_environment.py

@@ -315,7 +315,6 @@ def get_environment(self, step):
         # 5. All resolved packages (Conda or PyPI) are cached
         # 6. PyPI packages are only installed for local platform
 
-        # Resolve `linux-64` Conda environments if @batch or @kubernetes are in play
         target_platform = conda_platform()
         for decorator in step.decorators:
             # NOTE: Keep the list of supported decorator names for backward compatibility purposes.
@@ -329,7 +328,6 @@ def get_environment(self, step):
                 "snowpark",
                 "slurm",
             ]:
-                # TODO: Support arm architectures
                 target_platform = getattr(decorator, "target_platform", "linux-64")
                 break