Fix issues with experiment start and status check
@@ -5,21 +5,20 @@ import logging
from django.conf import settings
-from polyaxon_k8s.constants import ExperimentLifeCycle
from polyaxon_schemas.utils import TaskType
-from spawner import K8SSpawner
from api.celery_api import app as celery_app
from api.settings import CeleryTasks
-from experiments.models import ExperimentJob
from libs.redis_db import RedisExperimentStatus
+from spawner import K8SSpawner
+from spawner.utils.constants import ExperimentLifeCycle
logger = logging.getLogger('polyaxon.tasks.experiments')
@celery_app.task(name=CeleryTasks.EXPERIMENTS_START)
def start_experiment(experiment_id):
- from experiments.models import Experiment, ExperimentStatus
+ from experiments.models import Experiment, ExperimentStatus, ExperimentJob
try:
experiment = Experiment.objects.get(id=experiment_id)
@@ -31,19 +30,37 @@ def start_experiment(experiment_id):
ExperimentStatus.objects.create(experiment=experiment, status=ExperimentLifeCycle.SCHEDULED)
# Use spawner to start the experiment
- spawner = K8SSpawner(project_id=experiment.project.id,
- spec_id=experiment.spec.id if experiment.spec else '',
- experiment_id=experiment.id,
+ spawner = K8SSpawner(project_uuid=experiment.project.uuid.hex,
+ spec_uuid=experiment.spec.uuid.hex if experiment.spec else '',
+ experiment_uuid=experiment.uuid.hex,
specification=experiment.config,
- k8s_config=settings.K8S_CONFIG)
+ k8s_config=settings.K8S_CONFIG,
+ namespace=settings.K8S_NAMESPACE,
+ in_cluster=True,
+ use_sidecar=True,
+ sidecar_config='') # Use current settings config
resp = spawner.start_experiment()
# Get the number of jobs this experiment started
- ExperimentJob.objects.create(experiment=experiment, definition=resp[TaskType.MASTER])
+ master = resp[TaskType.MASTER]
+ job_uuid = master['labels']['task_id']
+ ExperimentJob.objects.create(uuid=job_uuid, experiment=experiment, definition=master)
for worker in resp[TaskType.WORKER]:
- ExperimentJob.objects.create(experiment=experiment, definition=worker)
+ job_uuid = worker['labels']['task_id']
+ ExperimentJob.objects.create(uuid=job_uuid, experiment=experiment, definition=worker)
for ps in resp[TaskType.PS]:
- ExperimentJob.objects.create(experiment=experiment, definition=ps)
+ job_uuid = ps['labels']['task_id']
+ ExperimentJob.objects.create(uuid=job_uuid, experiment=experiment, definition=ps)
# Add the experiment to the list of experiments to monitor
- RedisExperimentStatus.monitor(experiment_id)
+ RedisExperimentStatus.monitor(experiment.uuid.hex)
+
+
+@celery_app.task(name=CeleryTasks.EXPERIMENTS_CHECK_STATUS)
+def check_experiment_status(experiment_id):
+ from experiments.models import Experiment
+ from libs.redis_db import RedisExperimentStatus
+
+ experiment = Experiment.objects.get(id=experiment_id)
+ status = experiment.calculated_status
+ RedisExperimentStatus.set_status(experiment_uuid=experiment.uuid.hex, status=status)
