Fix extract_task_module

flytekit/core/tracker.py

@@ -298,12 +298,12 @@
     """
 
     if isinstance(f, TrackedInstance):
-        if f.instantiated_in:
+        if hasattr(f, "task_function"):
+            mod, mod_name, name = _task_module_from_callable(f.task_function)
+        elif f.instantiated_in:
             mod = importlib.import_module(f.instantiated_in)
             mod_name = mod.__name__
             name = f.lhs
-        elif hasattr(f, "task_function"):
-            mod, mod_name, name = _task_module_from_callable(f.task_function)
     else:
         mod, mod_name, name = _task_module_from_callable(f)
 

plugins/flytekit-papermill/tests/test_task.py

@@ -1,16 +1,23 @@
 import datetime
 import os
+import shutil
 import tempfile
 import typing
+from unittest import mock
 
 import pandas as pd
+from click.testing import CliRunner
 from flytekitplugins.papermill import NotebookTask
 from flytekitplugins.pod import Pod
 from kubernetes.client import V1Container, V1PodSpec
 
 import flytekit
 from flytekit import StructuredDataset, kwtypes, map_task, task, workflow
+from flytekit.clients.friendly import SynchronousFlyteClient
+from flytekit.clis.sdk_in_container import pyflyte
 from flytekit.configuration import Image, ImageConfig
+from flytekit.core import context_manager
+from flytekit.remote import FlyteRemote
 from flytekit.types.directory import FlyteDirectory
 from flytekit.types.file import FlyteFile, PythonNotebook
 
@@ -189,3 +196,32 @@ def wf(a: float) -> typing.List[float]:
         return map_task(nb_sub_task)(a=[a, a])
 
     assert wf(a=3.14) == [9.8596, 9.8596]
+
+
+@mock.patch("flytekit.clis.sdk_in_container.helpers.FlyteRemote", spec=FlyteRemote)
+@mock.patch("flytekit.clients.friendly.SynchronousFlyteClient", spec=SynchronousFlyteClient)
+def test_register_notebook_task(mock_client, mock_remote):
+    mock_remote._client = mock_client
+    mock_remote.return_value._version_from_hash.return_value = "dummy_version_from_hash"
+    mock_remote.return_value.fast_package.return_value = "dummy_md5_bytes", "dummy_native_url"
+    runner = CliRunner()
+    context_manager.FlyteEntities.entities.clear()
+    notebook_task = """
+from flytekitplugins.papermill import NotebookTask
+
+nb_simple = NotebookTask(
+    name="test",
+    notebook_path="./core/notebook.ipython",
+)
+"""
+    with runner.isolated_filesystem():
+        os.makedirs("core", exist_ok=True)
+        with open(os.path.join("core", "notebook.ipython"), "w") as f:
+            f.write("notebook.ipython")
+            f.close()
+        with open(os.path.join("core", "notebook_task.py"), "w") as f:
+            f.write(notebook_task)
+            f.close()
+        result = runner.invoke(pyflyte.main, ["register", "core"])
+        assert "Successfully registered 2 entities" in result.output
+        shutil.rmtree("core")

plugins/flytekit-sqlalchemy/tests/test_task.py

@@ -4,15 +4,21 @@
 import sqlite3
 import tempfile
 from typing import Iterator
+from unittest import mock
 
 import pandas
 import pytest
+from click.testing import CliRunner
 from flytekitplugins.sqlalchemy import SQLAlchemyConfig, SQLAlchemyTask
 from flytekitplugins.sqlalchemy.task import SQLAlchemyTaskExecutor
 
 from flytekit import kwtypes, task, workflow
+from flytekit.clients.friendly import SynchronousFlyteClient
+from flytekit.clis.sdk_in_container import pyflyte
+from flytekit.core import context_manager
 from flytekit.core.context_manager import SecretsManager
 from flytekit.models.security import Secret
+from flytekit.remote import FlyteRemote
 from flytekit.types.schema import FlyteSchema
 
 tk = SQLAlchemyTask(
@@ -197,3 +203,32 @@ def test_task_serialization_deserialization_with_secret(sql_server):
     r = executor.execute_from_model(tt)
 
     assert r.iat[0, 0] == 1
+
+
+@mock.patch("flytekit.clis.sdk_in_container.helpers.FlyteRemote", spec=FlyteRemote)
+@mock.patch("flytekit.clients.friendly.SynchronousFlyteClient", spec=SynchronousFlyteClient)
+def test_register_sql_task(mock_client, mock_remote):
+    mock_remote._client = mock_client
+    mock_remote.return_value._version_from_hash.return_value = "dummy_version_from_hash"
+    mock_remote.return_value.fast_package.return_value = "dummy_md5_bytes", "dummy_native_url"
+    runner = CliRunner()
+    context_manager.FlyteEntities.entities.clear()
+    sql_task = """
+from flytekitplugins.sqlalchemy import SQLAlchemyConfig, SQLAlchemyTask
+
+tk = SQLAlchemyTask(
+    "test",
+    query_template="select * from tracks",
+    task_config=SQLAlchemyConfig(
+        uri="sqlite://",
+    ),
+)
+"""
+    with runner.isolated_filesystem():
+        os.makedirs("core", exist_ok=True)
+        with open(os.path.join("core", "sql_task.py"), "w") as f:
+            f.write(sql_task)
+            f.close()
+        result = runner.invoke(pyflyte.main, ["register", "core"])
+        assert "Successfully registered 1 entities" in result.output
+        shutil.rmtree("core")

tests/flytekit/unit/cli/pyflyte/test_register.py

@@ -28,6 +28,15 @@ def my_workflow(x: int, y: int) -> int:
     return sum(x=square(z=x), y=square(z=y))
 """
 
+shell_task = """
+from flytekit.extras.tasks.shell import ShellTask
+
+t = ShellTask(
+        name="test",
+        script="echo 'Hello World'",
+    )
+"""
+
 
 @mock.patch("flytekit.clis.sdk_in_container.helpers.FlyteRemote")
 def test_saving_remote(mock_remote):
@@ -69,6 +78,26 @@ def test_register_with_no_output_dir_passed(mock_client, mock_remote):
         shutil.rmtree("core1")
 
 
+@mock.patch("flytekit.clis.sdk_in_container.helpers.FlyteRemote", spec=FlyteRemote)
+@mock.patch("flytekit.clients.friendly.SynchronousFlyteClient", spec=SynchronousFlyteClient)
+def test_register_shell_task(mock_client, mock_remote):
+    mock_remote._client = mock_client
+    mock_remote.return_value._version_from_hash.return_value = "dummy_version_from_hash"
+    mock_remote.return_value.fast_package.return_value = "dummy_md5_bytes", "dummy_native_url"
+    runner = CliRunner()
+    context_manager.FlyteEntities.entities.clear()
+    with runner.isolated_filesystem():
+        out = subprocess.run(["git", "init"], capture_output=True)
+        assert out.returncode == 0
+        os.makedirs("core2", exist_ok=True)
+        with open(os.path.join("core2", "shell_task.py"), "w") as f:
+            f.write(shell_task)
+            f.close()
+        result = runner.invoke(pyflyte.main, ["register", "core2"])
+        assert "Successfully registered 2 entities" in result.output
+        shutil.rmtree("core2")
+
+
 @mock.patch("flytekit.clis.sdk_in_container.helpers.FlyteRemote", spec=FlyteRemote)
 @mock.patch("flytekit.clients.friendly.SynchronousFlyteClient", spec=SynchronousFlyteClient)
 def test_non_fast_register(mock_client, mock_remote):