Groq autotracing #14006

anumita0203 · 2024-12-07T14:06:27Z

🛠 DevTools 🛠

Install mlflow from this PR

# Use `%sh` to run this command on Databricks
OPTIONS=$(if pip freeze | grep -q 'mlflow @ git+https://github.com/mlflow/mlflow.git'; then echo '--force-reinstall --no-deps'; fi)
pip install $OPTIONS git+https://github.com/mlflow/mlflow.git@refs/pull/14006/merge

Checkout with GitHub CLI

gh pr checkout 14006

Related Issues/PRs

Resolve #13911

What changes are proposed in this pull request?

Add Autolog tracing integration for Groq.

How is this PR tested?

Existing unit/integration tests
New unit/integration tests
Manual tests

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

No. You can skip the rest of this section.
Yes. Give a description of this change to be included in the release notes for MLflow users.

Add Autolog tracing integration for Groq.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

Interface

area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
area/docker: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
area/sqlalchemy: Use of SQLAlchemy in the Tracking Service or Model Registry
area/windows: Windows support

Language

language/r: R APIs and clients
language/java: Java APIs and clients
language/new: Proposals for new client languages

Integrations

integrations/azure: Azure and Azure ML integrations
integrations/sagemaker: SageMaker integrations
integrations/databricks: Databricks integrations

How should the PR be classified in the release notes? Choose one:

rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
rn/feature - A new user-facing feature worth mentioning in the release notes
rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
rn/documentation - A user-facing documentation change worth mentioning in the release notes

Should this PR be included in the next patch release?

Yes should be selected for bug fixes, documentation updates, and other small changes. No should be selected for new features and larger changes. If you're unsure about the release classification of this PR, leave this unchecked to let the maintainers decide.

What is a minor/patch release?

Minor release: a release that increments the second part of the version number (e.g., 1.2.0 -> 1.3.0).
Bug fixes, doc updates and new features usually go into minor releases.
Patch release: a release that increments the third part of the version number (e.g., 1.2.0 -> 1.2.1).
Bug fixes and doc updates usually go into patch releases.

Yes (this PR will be cherry-picked and included in the next patch release)
No (this PR will be included in the next minor release)

Signed-off-by: Anumita Dasgupta <anumita0302@gmail.com>

github-actions · 2024-12-07T14:06:41Z

@anumita0203 Thank you for the contribution! Could you fix the following issue(s)?

⚠ DCO check

The DCO check failed. Please sign off your commit(s) by following the instructions here. See https://github.com/mlflow/mlflow/blob/master/CONTRIBUTING.md#sign-your-work for more details.

github-actions · 2024-12-07T14:06:58Z

Documentation preview for bb3263a will be available when this CircleCI job
completes successfully.

More info

Ignore this comment if this PR does not change the documentation.
It takes a few minutes for the preview to be available.
The preview is updated when a new commit is pushed to this PR.
This comment was created by https://github.com/mlflow/mlflow/actions/runs/12497454459.

B-Step62 · 2024-12-09T02:27:56Z

mlflow/groq/__init__.py

Aside from the documentation, could you update a few more files to make CI pass? Reference: https://github.com/mlflow/mlflow/pull/13851/files

Add tests/groq to the conftest.py

Add groq to the ml-package-versions.yml file. The file defines supported versions and some installation commands for the CI. You can follow the format of anthropic. For maximum version, please set the latest version. Minimum version doesn't have to be very accurate.

B-Step62 · 2024-12-09T02:28:10Z

mlflow/groq/__init__.py

+from mlflow.groq._groq_autolog import (
+    patched_call,
+)


Suggested change

from mlflow.groq._groq_autolog import (

patched_call,

)

from mlflow.groq._groq_autolog import patched_call

B-Step62 · 2024-12-09T02:28:52Z

mlflow/groq/__init__.py

+    if Version(_get_groq_package_version()) < Version("0.13.0"):
+        raise MlflowException("Groq autologging is only supported for groq >= 0.13.0")


Once you added the minimum version in ml-package-versions.yml, you can remove this validation🙂

B-Step62 · 2024-12-09T02:30:32Z

mlflow/groq/_groq_autolog.py

+        ChatCompletions: SpanType.CHAT_MODEL,
+        Embeddings: SpanType.EMBEDDING,
+    }
+    return span_type_mapping.get(task, SpanType.UNKNOWN)


Suggested change

return span_type_mapping.get(task, SpanType.UNKNOWN)

return span_type_mapping.get(task, SpanType.LLM)

Let's use LLM type for translation / translations

B-Step62 · 2024-12-09T02:32:00Z

mlflow/groq/_groq_autolog.py

+    config = AutoLoggingConfig.init(flavor_name=mlflow.groq.FLAVOR_NAME)
+    active_run = mlflow.active_run()
+
+    # Active run should always take precedence over the run_id stored in the model
+    run_id = active_run.info.run_id if active_run else getattr(self, "_mlflow_run_id", None)


Only trace logging is required for now, so we don't need to manage Run logic. You can refer to anthropic implementation to simplify the logic here. (Sorry I should have refer it instead of OpenAI in the guidance).

B-Step62 · 2024-12-09T02:33:04Z

tests/groq/test_groq_autolog.py

Great test coverage!

Thanks 😊

Signed-off-by: Anumita Dasgupta <anumita0302@gmail.com>

B-Step62

The change overall looks great! I think we are very close to the goal🙂

A few more formal changes to make things work and test pass:

Add groq = LazyLoader("mlflow.groq", globals(), "mlflow.groq") in mlflow/init.py to allow lazy import.
Add "groq": "mlflow.groq" to this dictionary so that it is integrated to the global mlflow.autolog() API.

Once these are addressed, we should be good to merge and release this in next month. Thank you for putting effort on this, @anumita0203 !!

B-Step62 · 2024-12-18T00:06:43Z

tests/groq/test_groq_autolog.py

+    return DUMMY_CHAT_COMPLETION_RESPONSE
+
+
+def test_chat_completion_autolog():


According to the error message here, it seems we need to set some dummy API key in the environment variable for testing, even if we patch the client.

B-Step62 · 2024-12-18T00:09:28Z

mlflow/groq/__init__.py

+def _get_groq_package_version():
+    return importlib.metadata.version("groq")


nit: we can remove this function

Signed-off-by: Anumita Dasgupta <anumita0302@gmail.com>

Signed-off-by: Yuki Watanabe <31463517+B-Step62@users.noreply.github.com>

Signed-off-by: B-Step62 <yuki.watanabe@databricks.com>

B-Step62

LGTM!

Signed-off-by: Anumita Dasgupta <anumita0302@gmail.com> Signed-off-by: Yuki Watanabe <31463517+B-Step62@users.noreply.github.com> Signed-off-by: B-Step62 <yuki.watanabe@databricks.com> Co-authored-by: Yuki Watanabe <31463517+B-Step62@users.noreply.github.com> Co-authored-by: B-Step62 <yuki.watanabe@databricks.com> Signed-off-by: k99kurella <karthikkurella4@gmail.com>

anumita0203 added 4 commits December 6, 2024 17:26

introduce groq tracing

b949f22

Signed-off-by: Anumita Dasgupta <anumita0302@gmail.com>

add test cases

ac5faf0

Signed-off-by: Anumita Dasgupta <anumita0302@gmail.com>

Merge branch 'mlflow:master' into groq-autotracing

339f2f4

fix typo

e0a91d5

Signed-off-by: Anumita Dasgupta <anumita0302@gmail.com>

github-actions bot added the rn/feature Mention under Features in Changelogs. label Dec 7, 2024

anumita0203 mentioned this pull request Dec 7, 2024

[FR] MLflow Tracing for Groq #13911

Closed

B-Step62 reviewed Dec 9, 2024

View reviewed changes

anumita0203 added 3 commits December 12, 2024 17:50

address review comments

5f7678a

Signed-off-by: Anumita Dasgupta <anumita0302@gmail.com>

add missing documentation

74c75c3

Signed-off-by: Anumita Dasgupta <anumita0302@gmail.com>

fix typo

56a38db

Signed-off-by: Anumita Dasgupta <anumita0302@gmail.com>

B-Step62 reviewed Dec 18, 2024

View reviewed changes

B-Step62 mentioned this pull request Dec 18, 2024

Adds blog on patching a Groq client instance to add tracing support + some general principles on tracing customization mlflow/mlflow-website#135

Merged

anumita0203 and others added 4 commits December 25, 2024 13:06

minor changes

a7cf2fb

Signed-off-by: Anumita Dasgupta <anumita0302@gmail.com>

fix typo

09aaece

Signed-off-by: Anumita Dasgupta <anumita0302@gmail.com>

Merge branch 'master' into groq-autotracing

54d8d29

Signed-off-by: Yuki Watanabe <31463517+B-Step62@users.noreply.github.com>

fix

bb3263a

Signed-off-by: B-Step62 <yuki.watanabe@databricks.com>

B-Step62 approved these changes Dec 26, 2024

View reviewed changes

B-Step62 added this pull request to the merge queue Dec 26, 2024

Merged via the queue into mlflow:master with commit 7f056ad Dec 26, 2024
42 checks passed

harupy mentioned this pull request Jan 23, 2025

Run python3 dev/update_changelog.py --prev-version 2.2... #14314

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Groq autotracing #14006

Groq autotracing #14006

anumita0203 commented Dec 7, 2024 •

edited by github-actions bot

Loading

github-actions bot commented Dec 7, 2024

github-actions bot commented Dec 7, 2024 •

edited

Loading

B-Step62 Dec 9, 2024

B-Step62 Dec 9, 2024

B-Step62 Dec 9, 2024

B-Step62 Dec 9, 2024

B-Step62 Dec 9, 2024

B-Step62 Dec 9, 2024

anumita0203 Dec 12, 2024

B-Step62 left a comment

B-Step62 Dec 18, 2024

B-Step62 Dec 18, 2024

B-Step62 left a comment

		if Version(_get_groq_package_version()) < Version("0.13.0"):
		raise MlflowException("Groq autologging is only supported for groq >= 0.13.0")

	return span_type_mapping.get(task, SpanType.UNKNOWN)
	return span_type_mapping.get(task, SpanType.LLM)

		return DUMMY_CHAT_COMPLETION_RESPONSE


		def test_chat_completion_autolog():

		def _get_groq_package_version():
		return importlib.metadata.version("groq")

Groq autotracing #14006

Groq autotracing #14006

Conversation

anumita0203 commented Dec 7, 2024 • edited by github-actions bot Loading

Install mlflow from this PR

Checkout with GitHub CLI

Related Issues/PRs

What changes are proposed in this pull request?

How is this PR tested?

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

What component(s), interfaces, languages, and integrations does this PR affect?

How should the PR be classified in the release notes? Choose one:

Should this PR be included in the next patch release?

github-actions bot commented Dec 7, 2024

⚠ DCO check

github-actions bot commented Dec 7, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

B-Step62 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

B-Step62 left a comment

Choose a reason for hiding this comment

anumita0203 commented Dec 7, 2024 •

edited by github-actions bot

Loading

github-actions bot commented Dec 7, 2024 •

edited

Loading