Skip to content

Commit

Permalink
Remove local execution
Browse files Browse the repository at this point in the history
Signed-off-by: Edwin Yu <edwinyyyu@gmail.com>
  • Loading branch information
edwinyyyu committed Sep 15, 2023
1 parent e0edbe3 commit 7a1ab43
Showing 1 changed file with 2 additions and 10 deletions.
12 changes: 2 additions & 10 deletions plugins/flytekit-mmcloud/flytekitplugins/mmcloud/task.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
import os
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, Dict, Optional, Union
Expand All @@ -11,7 +10,6 @@
from flytekit.core.python_function_task import PythonFunctionTask
from flytekit.core.resources import Resources
from flytekit.extend import TaskPlugins
from flytekit.extend.backend.base_agent import AsyncAgentExecutorMixin
from flytekit.image_spec.image_spec import ImageSpec


Expand All @@ -25,7 +23,7 @@ class MMCloudConfig(object):
submit_extra: str = ""


class MMCloudTask(AsyncAgentExecutorMixin, PythonFunctionTask):
class MMCloudTask(PythonFunctionTask):
_TASK_TYPE = "mmcloud_task"

def __init__(
Expand All @@ -48,13 +46,7 @@ def __init__(
self._mmcloud_resources = flyte_to_mmcloud_resources(requests=requests, limits=limits)

def execute(self, **kwargs) -> Any:
# FLOAT_JOB_ID should always and only be defined on a Memory Machine Cloud worker node
if os.getenv("FLOAT_JOB_ID"):
# Task should be run as a normal PythonFunctionTask
return PythonFunctionTask.execute(self, **kwargs)
else:
# Assume local execution without a FlytePropeller deployment
return AsyncAgentExecutorMixin.execute(self, **kwargs)
return PythonFunctionTask.execute(self, **kwargs)

def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]:
"""
Expand Down

0 comments on commit 7a1ab43

Please sign in to comment.