WIP: Disable slow imports

it.py

@@ -0,0 +1,8 @@
+import os
+
+import django
+from django.urls import resolve
+
+os.environ.setdefault("UV_ENV_FILE", "./dev/.env.docker-compose")
+django.setup()
+resolve("/")

mixtape/core/models/clustering_result.py

@@ -5,7 +5,7 @@
 from django.contrib.postgres.fields import ArrayField
 from django.db import models
 
-from mixtape.core.analysis.constants import DEFAULT_CLUSTERING_PARAMS, FEATURE_TYPES
+# from mixtape.core.analysis.constants import DEFAULT_CLUSTERING_PARAMS, FEATURE_TYPES
 
 from .episode import Episode
 

mixtape/core/models/step.py

@@ -26,6 +26,7 @@ class Meta:
     number = models.PositiveIntegerField()
     image = models.ImageField(null=True, blank=True)
 
+    # TODO: move to algorithms?
     @contextmanager
     @staticmethod
     def rgb_array_to_file(rgb_array: npt.NDArray, file_base_name: str) -> Generator[File]:

mixtape/core/tasks/clustering_tasks.py

@@ -13,8 +13,8 @@
 from django.db.models.query import QuerySet
 from django.utils import timezone
 
-from mixtape.core.analysis.clustering import cluster_episodes_all_features
-from mixtape.core.analysis.constants import DEFAULT_CLUSTERING_PARAMS, DEFAULT_FEATURE_TYPES
+# from mixtape.core.analysis.clustering import cluster_episodes_all_features
+# from mixtape.core.analysis.constants import DEFAULT_CLUSTERING_PARAMS, DEFAULT_FEATURE_TYPES
 from mixtape.core.json_encoder import CustomJSONEncoder
 from mixtape.core.models import ClusteringResult, Episode
 
@@ -134,7 +134,9 @@ def compute_single_episode_clustering(episode_id: int) -> dict[str, Any]:
     Returns:
         Dict with task status and clustering result ID
     """
+    # TODO: why is this delaying another task?
     async_result = compute_episode_clustering.delay([episode_id])
+    # TODO: why is this returning a result in a task?
     return {
         'status': 'queued',
         'parent_task_id': async_result.id,

mixtape/core/views.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from collections import defaultdict
 from itertools import accumulate
 import json
@@ -9,15 +11,15 @@
 from django.http import Http404, HttpRequest, HttpResponse, JsonResponse
 from django.shortcuts import get_object_or_404, render
 import numpy as np
-from sklearn import cluster, decomposition, pipeline, preprocessing
-import umap.umap_ as umap
-
-from mixtape.core.analysis.clustering import (
-    _cluster_episodes_by_feature,
-    _fetch_all_episode_features,
-    cluster_episodes_all_features,
-)
-from mixtape.core.analysis.constants import DEFAULT_CLUSTERING_PARAMS
+# from sklearn import cluster, decomposition, pipeline, preprocessing
+# import umap.umap_ as umap
+
+# from mixtape.core.analysis.clustering import (
+#     _cluster_episodes_by_feature,
+#     _fetch_all_episode_features,
+#     cluster_episodes_all_features,
+# )
+# from mixtape.core.analysis.constants import DEFAULT_CLUSTERING_PARAMS
 from mixtape.core.json_encoder import CustomJSONEncoder
 from mixtape.core.models import ActionMapping, ClusteringResult, Episode
 from mixtape.core.tasks.clustering_tasks import (

pyproject.toml

@@ -11,8 +11,6 @@ license-files = ["LICENSE", "NOTICE"]
 classifiers = ["Private :: Do Not Upload"]
 dependencies = [
   # Runtime dependencies, always needed
-  "ale-py",
-  "autorom[accept-rom-license]",
   "celery",
   "django[argon2]",
   "django-allauth",
@@ -28,22 +26,9 @@ dependencies = [
   "django-stubs-ext",
   "djangorestframework",
   "drf-yasg",
-  "einops",
-  "gymnasium[atari, other]",
-  # PettingZoo[butterfly] extra pulls in both pymunk and pygame, but the pypi version of PettingZoo
-  # is extremely outdated and depends on a version of pymunk that does not build and has no bdist
-  # wheel for our target platforms. So we manually specify those transitive dependency versions.
-  "pymunk==6.11.0",
-  "pygame>=2.3.0",
-  "PettingZoo>=1.24.3",
-  "Pillow",
   "psycopg[binary]",
-  "ray[default, rllib]==2.34.0",
   "rich",
   "scikit-learn",
-  "SuperSuit",
-  "torch",
-  "umap-learn",
   "whitenoise[brotli]",
   # Production-only
   "django-s3-file-field[s3]",
@@ -64,6 +49,23 @@ development = [
   "watchdog",
   "werkzeug",
 ]
+analysis = [
+  "ale-py",
+  "autorom[accept-rom-license]",
+  "einops",
+  "gymnasium[atari, other]",
+  # PettingZoo[butterfly] extra pulls in both pymunk and pygame, but the pypi version of PettingZoo
+  # is extremely outdated and depends on a version of pymunk that does not build and has no bdist
+  # wheel for our target platforms. So we manually specify those transitive dependency versions.
+  "pymunk==6.11.0",
+  "pygame>=2.3.0",
+  "PettingZoo>=1.24.3",
+  "Pillow",
+  "ray[default, rllib]==2.34.0",
+  "SuperSuit",
+  "torch",
+  "umap-learn",
+]
 
 [dependency-groups]
 dev = [

tox.ini

@@ -13,6 +13,7 @@ runner = uv-venv-lock-runner
 pass_env =
     DJANGO_*
 extras =
+    analysis
     development
 
 [testenv:lint]